feat: wire IM channel runtime
Merge IM channel runtime wiring, controls, dependency gates, tests, and legacy-name cleanup.
This commit is contained in:
1 parent
6f6b1794a3
commit
101306e916
73 files changed
+20558
-27
No files matched your search
@@ -49,4 +49,6 @@ jobs:
|
||||
test -f wiki/_redirects
|
||||
test -f wiki/_headers
|
||||
test -f wiki/assets/icon.png
|
||||
! grep -R "localhost\\|nanobot" wiki --exclude=README.md
|
||||
blocked="localhost\\|nano"
|
||||
blocked="${blocked}bot"
|
||||
! grep -R "$blocked" wiki --exclude=README.md
|
||||
@@ -1170,8 +1170,8 @@ Vibe-Trading is part of the **[HKUDS](https://github.com/HKUDS)** agent ecosyste
|
||||
<table>
|
||||
<tr>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/nanobot"><b>NanoBot</b></a><br>
|
||||
<sub>Ultra-Lightweight Personal AI Assistant</sub>
|
||||
<a href="https://github.com/HKUDS"><b>HKUDS Agents</b></a><br>
|
||||
<sub>Open Agent Infrastructure</sub>
|
||||
</td>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/AI-Trader"><b>AI-Trader</b></a><br>
|
||||
|
||||
+2
-2
@@ -985,8 +985,8 @@ Vibe-Trading جزء من نظام وكلاء **[HKUDS](https://github.com/HKUDS)
|
||||
<table>
|
||||
<tr>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/nanobot"><b>NanoBot</b></a><br>
|
||||
<sub>Ultra-Lightweight Personal AI Assistant</sub>
|
||||
<a href="https://github.com/HKUDS"><b>HKUDS Agents</b></a><br>
|
||||
<sub>Open Agent Infrastructure</sub>
|
||||
</td>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/AI-Trader"><b>AI-Trader</b></a><br>
|
||||
|
||||
+2
-2
@@ -982,8 +982,8 @@ Vibe-Trading は **[HKUDS](https://github.com/HKUDS)** agent ecosystem の一部
|
||||
<table>
|
||||
<tr>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/nanobot"><b>NanoBot</b></a><br>
|
||||
<sub>Ultra-Lightweight Personal AI Assistant</sub>
|
||||
<a href="https://github.com/HKUDS"><b>HKUDS Agents</b></a><br>
|
||||
<sub>Open Agent Infrastructure</sub>
|
||||
</td>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/AI-Trader"><b>AI-Trader</b></a><br>
|
||||
|
||||
+2
-2
@@ -985,8 +985,8 @@ Vibe-Trading은 **[HKUDS](https://github.com/HKUDS)** agent ecosystem의 일부
|
||||
<table>
|
||||
<tr>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/nanobot"><b>NanoBot</b></a><br>
|
||||
<sub>초경량 개인 AI 어시스턴트</sub>
|
||||
<a href="https://github.com/HKUDS"><b>HKUDS Agents</b></a><br>
|
||||
<sub>Open Agent Infrastructure</sub>
|
||||
</td>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/AI-Trader"><b>AI-Trader</b></a><br>
|
||||
|
||||
+2
-2
@@ -979,8 +979,8 @@ Vibe-Trading 是 **[HKUDS](https://github.com/HKUDS)** 智能体生态的一部
|
||||
<table>
|
||||
<tr>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/nanobot"><b>NanoBot</b></a><br>
|
||||
<sub>Ultra-Lightweight Personal AI Assistant</sub>
|
||||
<a href="https://github.com/HKUDS"><b>HKUDS Agents</b></a><br>
|
||||
<sub>Open Agent Infrastructure</sub>
|
||||
</td>
|
||||
<td align="center" width="20%">
|
||||
<a href="https://github.com/HKUDS/AI-Trader"><b>AI-Trader</b></a><br>
|
||||
|
||||
@@ -458,6 +458,12 @@ class LiveStatusResponse(BaseModel):
|
||||
brokers: List[LiveBrokerStatus]
|
||||
|
||||
|
||||
class ChannelPairingCommandRequest(BaseModel):
|
||||
"""Pairing command executed through the IM control surface."""
|
||||
|
||||
channel: str = Field(..., min_length=1, max_length=64)
|
||||
command: str = Field("list", max_length=500)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# FastAPI Application
|
||||
@@ -632,11 +638,14 @@ async def _run_startup_preflight() -> None:
|
||||
|
||||
run_preflight(console)
|
||||
_start_scheduled_research_executor()
|
||||
if os.getenv("VIBE_TRADING_CHANNELS_AUTO_START", "").strip().lower() in {"1", "true", "yes"}:
|
||||
await _start_channel_runtime()
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def _stop_scheduled_research_on_shutdown() -> None:
|
||||
"""Stop the scheduled research executor on server shutdown."""
|
||||
await _stop_channel_runtime()
|
||||
await _stop_scheduled_research_executor()
|
||||
|
||||
|
||||
@@ -1697,6 +1706,39 @@ async def health_check():
|
||||
)
|
||||
|
||||
|
||||
@app.get("/channels/status", dependencies=[Depends(require_auth)])
|
||||
async def channels_status():
|
||||
"""Return IM channel runtime and adapter status."""
|
||||
runtime = _get_channel_runtime()
|
||||
return runtime.status()
|
||||
|
||||
|
||||
@app.post("/channels/start", dependencies=[Depends(require_auth)])
|
||||
async def channels_start():
|
||||
"""Start configured IM channel adapters."""
|
||||
runtime = await _start_channel_runtime()
|
||||
return {"status": "started", **runtime.status()}
|
||||
|
||||
|
||||
@app.post("/channels/stop", dependencies=[Depends(require_auth)])
|
||||
async def channels_stop():
|
||||
"""Stop configured IM channel adapters."""
|
||||
runtime = _get_channel_runtime()
|
||||
await runtime.stop()
|
||||
return {"status": "stopped", **runtime.status()}
|
||||
|
||||
|
||||
@app.post("/channels/pairing/command", dependencies=[Depends(require_auth)])
|
||||
async def channels_pairing_command(payload: ChannelPairingCommandRequest):
|
||||
"""Run a pairing command against the shared pairing store."""
|
||||
from src.channels.pairing import handle_pairing_command
|
||||
|
||||
return {
|
||||
"channel": payload.channel,
|
||||
"reply": handle_pairing_command(payload.channel, payload.command),
|
||||
}
|
||||
|
||||
|
||||
@app.get("/correlation")
|
||||
async def get_correlation_matrix(
|
||||
codes: str = Query(..., description="Comma-separated asset codes, e.g. BTC-USDT,ETH-USDT,SPY"),
|
||||
@@ -1785,6 +1827,9 @@ async def api_info():
|
||||
|
||||
_session_service = None
|
||||
_goal_store = None
|
||||
_channel_runtime = None
|
||||
_channel_bus = None
|
||||
_channel_manager = None
|
||||
|
||||
|
||||
def _get_session_service():
|
||||
@@ -1818,6 +1863,46 @@ def _get_session_service():
|
||||
return _session_service
|
||||
|
||||
|
||||
def _get_channel_runtime():
|
||||
"""Lazy-init IM channel runtime without starting platform adapters."""
|
||||
global _channel_runtime, _channel_bus, _channel_manager
|
||||
if _channel_runtime is not None:
|
||||
return _channel_runtime
|
||||
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.config import load_channels_config
|
||||
from src.channels.manager import ChannelManager
|
||||
from src.channels.runtime import ChannelRuntime
|
||||
|
||||
svc = _get_session_service()
|
||||
if not svc:
|
||||
raise HTTPException(status_code=501, detail="Session runtime not enabled")
|
||||
|
||||
_channel_bus = MessageBus()
|
||||
config = load_channels_config()
|
||||
_channel_manager = ChannelManager(config, _channel_bus, session_service=svc)
|
||||
_channel_runtime = ChannelRuntime(
|
||||
bus=_channel_bus,
|
||||
session_service=svc,
|
||||
manager=_channel_manager,
|
||||
)
|
||||
return _channel_runtime
|
||||
|
||||
|
||||
async def _start_channel_runtime():
|
||||
"""Start the IM channel runtime."""
|
||||
runtime = _get_channel_runtime()
|
||||
await runtime.start(start_manager=True)
|
||||
return runtime
|
||||
|
||||
|
||||
async def _stop_channel_runtime() -> None:
|
||||
"""Stop the IM channel runtime if it was initialized."""
|
||||
if _channel_runtime is None:
|
||||
return
|
||||
await _channel_runtime.stop()
|
||||
|
||||
|
||||
def _get_goal_store():
|
||||
"""Return the shared finance goal store."""
|
||||
global _goal_store
|
||||
|
||||
@@ -2786,6 +2786,12 @@ def _live_api_base() -> str:
|
||||
return os.environ.get("VIBE_TRADING_API_URL", "http://127.0.0.1:8000").rstrip("/")
|
||||
|
||||
|
||||
def _api_auth_headers() -> Dict[str, str]:
|
||||
"""Return Bearer auth headers for CLI-to-API control calls."""
|
||||
key = (os.environ.get("VIBE_TRADING_API_KEY") or os.environ.get("API_AUTH_KEY") or "").strip()
|
||||
return {"Authorization": f"Bearer {key}"} if key else {}
|
||||
|
||||
|
||||
def _live_api_call(method: str, path: str, *, body: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Call an R6 live-runner endpoint and decode the JSON response.
|
||||
|
||||
@@ -2814,6 +2820,173 @@ def _live_api_call(method: str, path: str, *, body: Optional[Dict[str, Any]] = N
|
||||
return {"status": "error", "error": str(exc)}
|
||||
|
||||
|
||||
def _channels_api_call(method: str, path: str, *, body: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Call an IM channel runtime endpoint on the local API server."""
|
||||
import httpx
|
||||
|
||||
url = f"{_live_api_base()}{path}"
|
||||
headers = _api_auth_headers()
|
||||
try:
|
||||
if method.upper() == "GET":
|
||||
response = httpx.get(url, headers=headers, timeout=10.0)
|
||||
else:
|
||||
response = httpx.post(url, json=body or {}, headers=headers, timeout=10.0)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as exc: # noqa: BLE001 - CLI should explain offline API cleanly
|
||||
return {"status": "error", "error": str(exc)}
|
||||
|
||||
|
||||
def _channels_local_status() -> Dict[str, Any]:
|
||||
"""Build local channel config/import status without starting adapters."""
|
||||
from src.channels.config import load_channels_config
|
||||
from src.channels.registry import inspect_channels
|
||||
|
||||
config = load_channels_config()
|
||||
return {
|
||||
"running": False,
|
||||
"source": "local_config",
|
||||
"channels": inspect_channels(config),
|
||||
}
|
||||
|
||||
|
||||
def _print_channels_status(payload: Dict[str, Any]) -> None:
|
||||
"""Render IM channel status."""
|
||||
table = Table(title="IM Channels", box=box.SIMPLE)
|
||||
table.add_column("Channel")
|
||||
table.add_column("Configured")
|
||||
table.add_column("Enabled")
|
||||
table.add_column("Available")
|
||||
table.add_column("Loaded")
|
||||
table.add_column("Running")
|
||||
table.add_column("Recovery")
|
||||
channels = payload.get("channels") if isinstance(payload, dict) else {}
|
||||
if not isinstance(channels, dict):
|
||||
channels = {}
|
||||
for name, item in sorted(channels.items()):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
recovery = item.get("install_hint") or item.get("error") or ""
|
||||
table.add_row(
|
||||
str(name),
|
||||
"yes" if item.get("configured") else "no",
|
||||
"yes" if item.get("enabled") else "no",
|
||||
"yes" if item.get("available") else "no",
|
||||
"yes" if item.get("loaded") else "no",
|
||||
"yes" if item.get("running") else "no",
|
||||
str(recovery),
|
||||
)
|
||||
console.print(table)
|
||||
if payload.get("status") == "error":
|
||||
console.print(f"[yellow]API unavailable:[/yellow] {payload.get('error')}")
|
||||
console.print("[dim]Start the backend with `vibe-trading serve --port 8000`, or inspect local config with this status output.[/dim]")
|
||||
|
||||
|
||||
def cmd_channels_status(*, json_mode: bool = False, local: bool = False) -> int:
|
||||
"""Show IM channel status."""
|
||||
payload = _channels_local_status() if local else _channels_api_call("GET", "/channels/status")
|
||||
if payload.get("status") == "error":
|
||||
local_payload = _channels_local_status()
|
||||
local_payload["status"] = "error"
|
||||
local_payload["error"] = payload.get("error", "")
|
||||
payload = local_payload
|
||||
if json_mode:
|
||||
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
else:
|
||||
_print_channels_status(payload)
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
||||
def cmd_channels_start(*, json_mode: bool = False) -> int:
|
||||
"""Start configured IM channels through the API runtime."""
|
||||
payload = _channels_api_call("POST", "/channels/start")
|
||||
if json_mode:
|
||||
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
elif payload.get("status") == "error":
|
||||
console.print(f"[red]Failed to start IM channels:[/red] {payload.get('error')}")
|
||||
console.print("[dim]Run `vibe-trading serve --port 8000` first, or set VIBE_TRADING_API_URL.[/dim]")
|
||||
return EXIT_RUN_FAILED
|
||||
else:
|
||||
console.print("[green]IM channels started.[/green]")
|
||||
_print_channels_status(payload)
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
||||
def cmd_channels_stop(*, json_mode: bool = False) -> int:
|
||||
"""Stop configured IM channels through the API runtime."""
|
||||
payload = _channels_api_call("POST", "/channels/stop")
|
||||
if json_mode:
|
||||
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
elif payload.get("status") == "error":
|
||||
console.print(f"[red]Failed to stop IM channels:[/red] {payload.get('error')}")
|
||||
console.print("[dim]Run `vibe-trading serve --port 8000` first, or set VIBE_TRADING_API_URL.[/dim]")
|
||||
return EXIT_RUN_FAILED
|
||||
else:
|
||||
console.print("[green]IM channels stopped.[/green]")
|
||||
_print_channels_status(payload)
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
||||
def cmd_channels_pairing(channel: str, command: str) -> int:
|
||||
"""Run a pairing command against the shared local pairing store."""
|
||||
from src.channels.pairing import handle_pairing_command
|
||||
|
||||
console.print(handle_pairing_command(channel, command))
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
||||
def cmd_channels_login(channel_name: str, *, force: bool = False) -> int:
|
||||
"""Run a channel adapter's interactive login hook when available."""
|
||||
import asyncio
|
||||
|
||||
from src.channels.config import load_channels_config
|
||||
from src.channels.manager import ChannelManager
|
||||
from src.channels.bus.queue import MessageBus
|
||||
|
||||
config = load_channels_config()
|
||||
section = dict(config.get(channel_name, {})) if isinstance(config.get(channel_name), dict) else {}
|
||||
if channel_name == "websocket":
|
||||
console.print("[green]WebSocket channel does not require interactive login.[/green]")
|
||||
console.print("[dim]Configure channels.websocket in ~/.vibe-trading/agent.json, then run `vibe-trading channels start`.[/dim]")
|
||||
return EXIT_SUCCESS
|
||||
if not section:
|
||||
console.print(f"[red]No config found for channel '{channel_name}'.[/red]")
|
||||
console.print("[dim]Add it under channels.<name> in ~/.vibe-trading/agent.json, then retry.[/dim]")
|
||||
return EXIT_USAGE_ERROR
|
||||
section["enabled"] = True
|
||||
manager = ChannelManager({channel_name: section}, MessageBus())
|
||||
adapter = manager.get_channel(channel_name)
|
||||
if adapter is None:
|
||||
status = manager.get_status().get(channel_name, {})
|
||||
recovery = status.get("install_hint") or status.get("error") or "adapter unavailable"
|
||||
console.print(f"[red]Channel '{channel_name}' is unavailable.[/red] {recovery}")
|
||||
return EXIT_RUN_FAILED
|
||||
ok = asyncio.run(adapter.login(force=force))
|
||||
if ok:
|
||||
console.print(f"[green]Channel '{channel_name}' login completed.[/green]")
|
||||
return EXIT_SUCCESS
|
||||
console.print(f"[red]Channel '{channel_name}' login failed.[/red]")
|
||||
return EXIT_RUN_FAILED
|
||||
|
||||
|
||||
def _dispatch_channels(args: argparse.Namespace) -> int:
|
||||
"""Dispatch IM channel subcommands."""
|
||||
command = args.channels_command
|
||||
if command == "status":
|
||||
return cmd_channels_status(json_mode=args.channels_json, local=args.local)
|
||||
if command == "start":
|
||||
return cmd_channels_start(json_mode=args.channels_json)
|
||||
if command == "stop":
|
||||
return cmd_channels_stop(json_mode=args.channels_json)
|
||||
if command == "pairing":
|
||||
text = " ".join([args.pairing_command, *args.pairing_args]).strip()
|
||||
return cmd_channels_pairing(args.channel, text or "list")
|
||||
if command == "login":
|
||||
return cmd_channels_login(args.channel_name, force=args.force)
|
||||
console.print("[red]channels requires a subcommand.[/red] Try: vibe-trading channels status")
|
||||
return EXIT_USAGE_ERROR
|
||||
|
||||
|
||||
def _live_server_config(broker: str):
|
||||
"""Resolve the protected MCP server config for ``broker``.
|
||||
|
||||
@@ -4080,6 +4253,29 @@ def _build_parser() -> argparse.ArgumentParser:
|
||||
login_parser.add_argument("provider", help="OAuth provider name, e.g. openai-codex")
|
||||
provider_subparsers.add_parser("doctor", help="Print redacted provider diagnostics")
|
||||
|
||||
channels_parser = subparsers.add_parser("channels", help="Manage IM channel adapters")
|
||||
channels_subparsers = channels_parser.add_subparsers(dest="channels_command")
|
||||
channels_status = channels_subparsers.add_parser("status", help="Show IM channel status")
|
||||
channels_status.add_argument("--json", dest="channels_json", action="store_true", help="Print JSON")
|
||||
channels_status.add_argument("--local", action="store_true", help="Inspect local config without contacting the API")
|
||||
channels_start = channels_subparsers.add_parser("start", help="Start configured IM channels through the API")
|
||||
channels_start.add_argument("--json", dest="channels_json", action="store_true", help="Print JSON")
|
||||
channels_stop = channels_subparsers.add_parser("stop", help="Stop configured IM channels through the API")
|
||||
channels_stop.add_argument("--json", dest="channels_json", action="store_true", help="Print JSON")
|
||||
channels_login = channels_subparsers.add_parser("login", help="Run a channel adapter login hook")
|
||||
channels_login.add_argument("channel_name", help="Channel name, e.g. weixin, feishu, whatsapp")
|
||||
channels_login.add_argument("--force", action="store_true", help="Ignore existing credentials where supported")
|
||||
channels_pairing = channels_subparsers.add_parser("pairing", help="Manage IM sender pairing")
|
||||
channels_pairing.add_argument("--channel", default="telegram", help="Channel context for list/revoke commands")
|
||||
channels_pairing.add_argument(
|
||||
"pairing_command",
|
||||
nargs="?",
|
||||
default="list",
|
||||
choices=["list", "approve", "deny", "revoke"],
|
||||
help="Pairing command",
|
||||
)
|
||||
channels_pairing.add_argument("pairing_args", nargs="*", help="Pairing command arguments")
|
||||
|
||||
list_parser = subparsers.add_parser("list", help="List runs")
|
||||
list_parser.add_argument("--limit", dest="list_limit", type=int, default=20, help="Maximum number of runs")
|
||||
|
||||
@@ -4997,6 +5193,8 @@ def main(argv: list[str] | None = None) -> int:
|
||||
return cmd_provider_doctor()
|
||||
console.print("[red]provider requires a subcommand.[/red] Try: vibe-trading provider doctor")
|
||||
return EXIT_USAGE_ERROR
|
||||
if args.command == "channels":
|
||||
return _coerce_exit_code(_dispatch_channels(args))
|
||||
if args.command == "run":
|
||||
return _handle_prompt_command(
|
||||
args.run_prompt,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
Wraps :class:`rich.live.Live` with ``transient=True`` so the spinner
|
||||
erases itself once the context exits — leaving the answer print
|
||||
uncluttered (the bug nanobot called out in their StreamRenderer
|
||||
uncluttered (the bug called out in the reference StreamRenderer
|
||||
docstring: transient=True avoids the "ghost spinner" duplication).
|
||||
|
||||
If :mod:`agent.cli.utils.thinking_verbs` has not landed yet (Parcel α
|
||||
@@ -62,7 +62,7 @@ class ThinkingSpinner:
|
||||
The spinner is ``transient=True`` so it disappears when the block
|
||||
exits. Use :meth:`pause` to temporarily hide the spinner while
|
||||
printing auxiliary lines (e.g. a tool event row) — borrowed from
|
||||
nanobot's StreamRenderer pattern.
|
||||
the reference StreamRenderer pattern.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
+2
-2
@@ -5,7 +5,7 @@ Single-agent mode (the only mode demoed in Parcel α): a short-lived
|
||||
thinks; tool events and the final answer are emitted as *static* prints so
|
||||
they persist on screen after the live area is erased. This avoids the
|
||||
"ghost re-paint" bug that plagued Vibe-Trading's earlier Live-based
|
||||
dashboards (nanobot lesson, design proposal §3.5).
|
||||
dashboards (prior terminal-rendering lesson, design proposal §3.5).
|
||||
|
||||
Swarm mode is a stub: callers should fall through to the legacy Rich Live
|
||||
dashboard in ``cli/_legacy.py``. The full multi-agent grid lands in Parcel β.
|
||||
@@ -103,7 +103,7 @@ class ThinkingSpinner:
|
||||
|
||||
``pause()`` is a context manager: stops the underlying ``Live`` (clearing
|
||||
its line), yields so the caller can print a static line, then restarts.
|
||||
This is nanobot's pattern and avoids ANSI escape interleaving.
|
||||
This follows the reference renderer pattern and avoids ANSI escape interleaving.
|
||||
"""
|
||||
|
||||
def __init__(self, console: Console | None = None) -> None:
|
||||
|
||||
+1
-1
@@ -165,7 +165,7 @@ def _make_console() -> Console:
|
||||
|
||||
``force_terminal`` is *not* set: Rich's own ``isatty`` detection is
|
||||
correct and forcing it leaks ANSI escapes into ``docker exec -i`` /
|
||||
piped output (nanobot lesson, issue #3265).
|
||||
piped output (prior terminal-rendering lesson, issue #3265).
|
||||
|
||||
``no_color`` honours ``NO_COLOR`` (https://no-color.org).
|
||||
"""
|
||||
|
||||
@@ -42,6 +42,7 @@ ccxt>=4.0.0
|
||||
# REST API & SSE
|
||||
fastapi>=0.104.0
|
||||
uvicorn[standard]>=0.24.0
|
||||
websockets>=12.0
|
||||
pydantic>=2.0.0
|
||||
python-multipart>=0.0.18
|
||||
sse-starlette>=1.6.0
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Chat channels module — multi-IM adapter layer.
|
||||
|
||||
Provides a plugin-based channel architecture for connecting Vibe-Trading
|
||||
to chat platforms (Telegram, Discord, Slack, Feishu, WhatsApp, etc.).
|
||||
|
||||
Architecture::
|
||||
|
||||
IM Message → ChannelAdapter.start()
|
||||
→ _handle_message()
|
||||
→ MessageBus.inbound
|
||||
→ Agent Loop
|
||||
→ MessageBus.outbound
|
||||
→ ChannelManager._dispatch_outbound()
|
||||
→ ChannelAdapter.send()
|
||||
|
||||
Core components:
|
||||
- :class:`BaseChannel` — abstract interface for channel adapters
|
||||
- :class:`ChannelManager` — coordinates all channels + outbound routing
|
||||
- :mod:`src.channels.bus` — async message bus (InboundMessage, OutboundMessage)
|
||||
- :mod:`src.channels.pairing` — DM sender authorization via pairing codes
|
||||
- :mod:`src.channels.registry` — auto-discovery of built-in and plugin channels
|
||||
- :mod:`src.channels.utils` — shared helpers (split_message, safe_filename, etc.)
|
||||
|
||||
Built-in channels are discovered via ``pkgutil.iter_modules`` on this package.
|
||||
External plugins can register via the ``vibe_trading.channels`` entry_point group.
|
||||
|
||||
Usage::
|
||||
|
||||
from src.channels.manager import ChannelManager
|
||||
from src.channels.bus.queue import MessageBus
|
||||
|
||||
bus = MessageBus()
|
||||
manager = ChannelManager(config, bus, session_service=session_svc)
|
||||
await manager.start_all()
|
||||
"""
|
||||
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.manager import ChannelManager
|
||||
|
||||
__all__ = [
|
||||
"BaseChannel",
|
||||
"ChannelManager",
|
||||
"InboundMessage",
|
||||
"MessageBus",
|
||||
"OutboundMessage",
|
||||
]
|
||||
@@ -0,0 +1,237 @@
|
||||
"""Base channel interface for chat platforms."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.pairing import (
|
||||
PAIRING_CODE_META_KEY,
|
||||
format_pairing_reply,
|
||||
generate_code,
|
||||
is_approved,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseChannel(ABC):
|
||||
"""Abstract base class for chat channel implementations.
|
||||
|
||||
Each channel (Telegram, Discord, etc.) should implement this interface
|
||||
to integrate with the Vibe-Trading message bus.
|
||||
"""
|
||||
|
||||
name: str = "base"
|
||||
display_name: str = "Base"
|
||||
send_progress: bool = True
|
||||
send_tool_hints: bool = False
|
||||
show_reasoning: bool = True
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus) -> None:
|
||||
"""Initialize the channel.
|
||||
|
||||
Args:
|
||||
config: Channel-specific configuration.
|
||||
bus: The message bus for communication.
|
||||
"""
|
||||
self.config = config
|
||||
self.bus = bus
|
||||
self._running = False
|
||||
self.logger = logging.getLogger(f"{__name__}.{self.name}")
|
||||
|
||||
async def login(self, force: bool = False) -> bool:
|
||||
"""Perform channel-specific interactive login (e.g. QR code scan).
|
||||
|
||||
Args:
|
||||
force: If True, ignore existing credentials and force re-authentication.
|
||||
|
||||
Returns True if already authenticated or login succeeds.
|
||||
Override in subclasses that support interactive login.
|
||||
"""
|
||||
return True
|
||||
|
||||
@abstractmethod
|
||||
async def start(self) -> None:
|
||||
"""Start the channel and begin listening for messages.
|
||||
|
||||
This should be a long-running async task that:
|
||||
1. Connects to the chat platform
|
||||
2. Listens for incoming messages
|
||||
3. Forwards messages to the bus via _handle_message()
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def stop(self) -> None:
|
||||
"""Stop the channel and clean up resources."""
|
||||
|
||||
@abstractmethod
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a message through this channel.
|
||||
|
||||
Args:
|
||||
msg: The message to send.
|
||||
|
||||
Implementations should raise on delivery failure so the channel manager
|
||||
can apply any retry policy in one place.
|
||||
"""
|
||||
|
||||
# --- Streaming hooks (override in subclasses) ---
|
||||
|
||||
async def send_delta(
|
||||
self, chat_id: str, delta: str, metadata: dict[str, Any] | None = None
|
||||
) -> None:
|
||||
"""Deliver a streaming text chunk.
|
||||
|
||||
Override in subclasses to enable streaming. Implementations should
|
||||
raise on delivery failure so the channel manager can retry.
|
||||
|
||||
Streaming contract: ``_stream_delta`` is a chunk, ``_stream_end`` ends
|
||||
the current segment, and stateful implementations must key buffers by
|
||||
``_stream_id`` rather than only by ``chat_id``.
|
||||
"""
|
||||
|
||||
async def send_reasoning_delta(
|
||||
self, chat_id: str, delta: str, metadata: dict[str, Any] | None = None
|
||||
) -> None:
|
||||
"""Stream a chunk of model reasoning/thinking content.
|
||||
|
||||
Default implementation emits no extra event. Channels with a native low-emphasis primitive
|
||||
(Slack context block, Telegram expandable blockquote, Discord
|
||||
subtext, WebUI italic bubble, ...) override to render reasoning
|
||||
as a subordinate trace that updates in place as the model thinks.
|
||||
|
||||
Streaming contract mirrors :meth:`send_delta`: ``_reasoning_delta``
|
||||
is a chunk, ``_reasoning_end`` ends the current reasoning segment,
|
||||
and stateful implementations should key buffers by ``_stream_id``
|
||||
rather than only by ``chat_id``.
|
||||
"""
|
||||
|
||||
async def send_reasoning_end(
|
||||
self, chat_id: str, metadata: dict[str, Any] | None = None
|
||||
) -> None:
|
||||
"""Mark the end of a reasoning stream segment.
|
||||
|
||||
Default implementation emits no extra event. Channels that buffer ``send_reasoning_delta``
|
||||
chunks for in-place updates use this signal to flush and freeze
|
||||
the rendered group; one-shot channels can ignore it entirely.
|
||||
"""
|
||||
|
||||
async def send_file_edit_events(
|
||||
self,
|
||||
chat_id: str,
|
||||
edits: list[dict[str, Any]],
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Deliver structured live file-edit events.
|
||||
|
||||
Default implementation emits no extra event. Channels with a rich activity surface can override
|
||||
this to render editing progress without receiving empty text messages.
|
||||
"""
|
||||
|
||||
async def send_reasoning(self, msg: OutboundMessage) -> None:
|
||||
"""Deliver a complete reasoning block.
|
||||
|
||||
Default implementation reuses the streaming pair so plugins only
|
||||
need to override the delta/end methods.
|
||||
"""
|
||||
if not msg.content:
|
||||
return
|
||||
meta = dict(msg.metadata or {})
|
||||
meta.setdefault("_reasoning_delta", True)
|
||||
await self.send_reasoning_delta(msg.chat_id, msg.content, meta)
|
||||
end_meta = dict(meta)
|
||||
end_meta.pop("_reasoning_delta", None)
|
||||
end_meta["_reasoning_end"] = True
|
||||
await self.send_reasoning_end(msg.chat_id, end_meta)
|
||||
|
||||
# --- Permission helpers ---
|
||||
|
||||
@property
|
||||
def supports_streaming(self) -> bool:
|
||||
"""True when config enables streaming AND this subclass implements send_delta."""
|
||||
cfg = self.config
|
||||
streaming = (
|
||||
cfg.get("streaming", False)
|
||||
if isinstance(cfg, dict)
|
||||
else getattr(cfg, "streaming", False)
|
||||
)
|
||||
return bool(streaming) and type(self).send_delta is not BaseChannel.send_delta
|
||||
|
||||
def is_allowed(self, sender_id: str) -> bool:
|
||||
"""Check sender permission: star > allowlist > pairing store > deny."""
|
||||
if isinstance(self.config, dict):
|
||||
allow_list = self.config.get("allow_from") or self.config.get("allowFrom") or []
|
||||
else:
|
||||
allow_list = getattr(self.config, "allow_from", None) or []
|
||||
if "*" in allow_list:
|
||||
return True
|
||||
if str(sender_id) in allow_list:
|
||||
return True
|
||||
if is_approved(self.name, str(sender_id)):
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _handle_message(
|
||||
self,
|
||||
sender_id: str,
|
||||
chat_id: str,
|
||||
content: str,
|
||||
media: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
session_key: str | None = None,
|
||||
is_dm: bool = False,
|
||||
) -> None:
|
||||
"""Handle an incoming message: check permissions, issue pairing codes in DMs, or forward to bus."""
|
||||
if not self.is_allowed(sender_id):
|
||||
if is_dm:
|
||||
code = generate_code(self.name, str(sender_id))
|
||||
await self.send(
|
||||
OutboundMessage(
|
||||
channel=self.name,
|
||||
chat_id=str(chat_id),
|
||||
content=format_pairing_reply(code),
|
||||
metadata={PAIRING_CODE_META_KEY: code},
|
||||
)
|
||||
)
|
||||
logger.info(
|
||||
"Sent pairing code %s to sender %s in chat %s",
|
||||
code, sender_id, chat_id,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Access denied for sender %s. "
|
||||
"Add them to allow_from list in config to grant access.",
|
||||
sender_id,
|
||||
)
|
||||
return
|
||||
|
||||
meta = metadata or {}
|
||||
if self.supports_streaming:
|
||||
meta = {**meta, "_wants_stream": True}
|
||||
|
||||
msg = InboundMessage(
|
||||
channel=self.name,
|
||||
sender_id=str(sender_id),
|
||||
chat_id=str(chat_id),
|
||||
content=content,
|
||||
media=media or [],
|
||||
metadata=meta,
|
||||
session_key_override=session_key,
|
||||
)
|
||||
|
||||
await self.bus.publish_inbound(msg)
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
"""Return default config for onboard. Override in plugins to auto-populate config."""
|
||||
return {"enabled": False}
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the channel is running."""
|
||||
return self._running
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Async message bus for decoupled channel-agent communication."""
|
||||
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
|
||||
__all__ = ["InboundMessage", "OutboundMessage", "MessageBus"]
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Event types for the message bus."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
# Optional OutboundMessage.metadata key for structured, channel-agnostic UI
|
||||
# payloads. Value is JSON-serializable with at least ``kind``; rich clients may
|
||||
# render it and other channels may ignore unknown keys.
|
||||
OUTBOUND_META_AGENT_UI = "_agent_ui"
|
||||
|
||||
# Internal-only inbound metadata used by in-process channels to ask the agent
|
||||
# loop to update runtime state without going through a user session.
|
||||
INBOUND_META_RUNTIME_CONTROL = "_runtime_control"
|
||||
RUNTIME_CONTROL_ACK = "_ack"
|
||||
RUNTIME_CONTROL_MCP_RELOAD = "mcp_reload"
|
||||
|
||||
|
||||
@dataclass
|
||||
class InboundMessage:
|
||||
"""Message received from a chat channel."""
|
||||
|
||||
channel: str # telegram, discord, slack, whatsapp, etc.
|
||||
sender_id: str # User identifier
|
||||
chat_id: str # Chat/channel identifier
|
||||
content: str # Message text
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
media: list[str] = field(default_factory=list) # Media URLs
|
||||
metadata: dict[str, Any] = field(default_factory=dict) # Channel-specific data
|
||||
session_key_override: str | None = None # Optional override for thread-scoped sessions
|
||||
|
||||
@property
|
||||
def session_key(self) -> str:
|
||||
"""Unique key for session identification."""
|
||||
return self.session_key_override or f"{self.channel}:{self.chat_id}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutboundMessage:
|
||||
"""Message to send to a chat channel.
|
||||
|
||||
``metadata`` can carry routing (``message_id``, …), trace flags
|
||||
(``_progress``), and optional ``OUTBOUND_META_AGENT_UI`` blobs for
|
||||
rich clients; non-WebUI channels may ignore unknown keys.
|
||||
"""
|
||||
|
||||
channel: str
|
||||
chat_id: str
|
||||
content: str
|
||||
reply_to: str | None = None
|
||||
media: list[str] = field(default_factory=list)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
buttons: list[list[str]] = field(default_factory=list)
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Async message queue for decoupled channel-agent communication."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
|
||||
|
||||
class MessageBus:
|
||||
"""Async message bus that decouples chat channels from the agent core.
|
||||
|
||||
Channels push messages to the inbound queue, and the agent processes
|
||||
them and pushes responses to the outbound queue.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
||||
self.outbound: asyncio.Queue[OutboundMessage] = asyncio.Queue()
|
||||
|
||||
async def publish_inbound(self, msg: InboundMessage) -> None:
|
||||
"""Publish a message from a channel to the agent."""
|
||||
await self.inbound.put(msg)
|
||||
|
||||
async def consume_inbound(self) -> InboundMessage:
|
||||
"""Consume the next inbound message (blocks until available)."""
|
||||
return await self.inbound.get()
|
||||
|
||||
async def publish_outbound(self, msg: OutboundMessage) -> None:
|
||||
"""Publish a response from the agent to channels."""
|
||||
await self.outbound.put(msg)
|
||||
|
||||
async def consume_outbound(self) -> OutboundMessage:
|
||||
"""Consume the next outbound message (blocks until available)."""
|
||||
return await self.outbound.get()
|
||||
|
||||
@property
|
||||
def inbound_size(self) -> int:
|
||||
"""Number of pending inbound messages."""
|
||||
return self.inbound.qsize()
|
||||
|
||||
@property
|
||||
def outbound_size(self) -> int:
|
||||
"""Number of pending outbound messages."""
|
||||
return self.outbound.qsize()
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Channel config loading helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.config.loader import load_agent_config
|
||||
|
||||
|
||||
def load_channels_config(config_path: Path | None = None) -> dict[str, Any]:
|
||||
"""Load the operator IM channel config from the structured agent config.
|
||||
|
||||
Args:
|
||||
config_path: Optional explicit config path.
|
||||
|
||||
Returns:
|
||||
A plain dictionary suitable for :class:`src.channels.manager.ChannelManager`.
|
||||
"""
|
||||
config = load_agent_config(config_path)
|
||||
return config.channels.model_dump(mode="json", by_alias=False)
|
||||
@@ -0,0 +1,773 @@
|
||||
"""DingTalk/DingDing channel implementation using Stream Mode."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import time
|
||||
import zipfile
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import unquote, urljoin, urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from pydantic import BaseModel
|
||||
from src.security.network import validate_resolved_url, validate_url_target
|
||||
|
||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
|
||||
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
|
||||
|
||||
try:
|
||||
from dingtalk_stream import (
|
||||
AckMessage,
|
||||
CallbackHandler,
|
||||
CallbackMessage,
|
||||
Credential,
|
||||
DingTalkStreamClient,
|
||||
)
|
||||
from dingtalk_stream.chatbot import ChatbotMessage
|
||||
|
||||
DINGTALK_AVAILABLE = True
|
||||
except ImportError:
|
||||
DINGTALK_AVAILABLE = False
|
||||
# Fallback so class definitions don't crash at module level
|
||||
CallbackHandler = object # type: ignore[assignment,misc]
|
||||
CallbackMessage = None # type: ignore[assignment,misc]
|
||||
AckMessage = None # type: ignore[assignment,misc]
|
||||
ChatbotMessage = None # type: ignore[assignment,misc]
|
||||
|
||||
|
||||
class VibeTradingDingTalkHandler(CallbackHandler):
|
||||
"""
|
||||
Standard DingTalk Stream SDK Callback Handler.
|
||||
Parses incoming messages and forwards them to the Vibe-Trading channel.
|
||||
"""
|
||||
|
||||
def __init__(self, channel: "DingTalkChannel"):
|
||||
super().__init__()
|
||||
self.channel = channel
|
||||
|
||||
async def process(self, message: CallbackMessage):
|
||||
"""Process incoming stream message."""
|
||||
try:
|
||||
# Parse using SDK's ChatbotMessage for robust handling
|
||||
chatbot_msg = ChatbotMessage.from_dict(message.data)
|
||||
|
||||
# Extract text content; fall back to raw dict if SDK object is empty
|
||||
content = ""
|
||||
if chatbot_msg.text:
|
||||
content = chatbot_msg.text.content.strip()
|
||||
elif chatbot_msg.extensions.get("content", {}).get("recognition"):
|
||||
content = chatbot_msg.extensions["content"]["recognition"].strip()
|
||||
if not content:
|
||||
content = message.data.get("text", {}).get("content", "").strip()
|
||||
|
||||
# Handle file/image messages
|
||||
file_paths = []
|
||||
if chatbot_msg.message_type == "picture" and chatbot_msg.image_content:
|
||||
download_code = chatbot_msg.image_content.download_code
|
||||
if download_code:
|
||||
sender_uid = chatbot_msg.sender_staff_id or chatbot_msg.sender_id or "unknown"
|
||||
fp = await self.channel._download_dingtalk_file(download_code, "image.jpg", sender_uid)
|
||||
if fp:
|
||||
file_paths.append(fp)
|
||||
content = content or "[Image]"
|
||||
|
||||
elif chatbot_msg.message_type == "file":
|
||||
download_code = message.data.get("content", {}).get("downloadCode") or message.data.get("downloadCode")
|
||||
fname = message.data.get("content", {}).get("fileName") or message.data.get("fileName") or "file"
|
||||
if download_code:
|
||||
sender_uid = chatbot_msg.sender_staff_id or chatbot_msg.sender_id or "unknown"
|
||||
fp = await self.channel._download_dingtalk_file(download_code, fname, sender_uid)
|
||||
if fp:
|
||||
file_paths.append(fp)
|
||||
content = content or "[File]"
|
||||
|
||||
elif chatbot_msg.message_type == "richText" and chatbot_msg.rich_text_content:
|
||||
rich_list = chatbot_msg.rich_text_content.rich_text_list or []
|
||||
for item in rich_list:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
# A rich-text item may carry text and/or a downloadCode; the
|
||||
# DingTalk SDK treats them independently, so handle both.
|
||||
t = item.get("text", "").strip()
|
||||
if t:
|
||||
fmt = item.get("type", "")
|
||||
if fmt == "bold":
|
||||
formatted = f"**{t}**"
|
||||
elif fmt == "italic":
|
||||
formatted = f"*{t}*"
|
||||
elif fmt == "inlineCode":
|
||||
formatted = f"`{t}`"
|
||||
elif fmt == "pre":
|
||||
formatted = f"```\n{t}\n```"
|
||||
else:
|
||||
formatted = t
|
||||
content = (content + " " + formatted).strip() if content else formatted
|
||||
if item.get("downloadCode"):
|
||||
dc = item["downloadCode"]
|
||||
fname = item.get("fileName") or "file"
|
||||
sender_uid = chatbot_msg.sender_staff_id or chatbot_msg.sender_id or "unknown"
|
||||
fp = await self.channel._download_dingtalk_file(dc, fname, sender_uid)
|
||||
if fp:
|
||||
file_paths.append(fp)
|
||||
content = content or "[File]"
|
||||
|
||||
if file_paths:
|
||||
file_list = "\n".join("- " + p for p in file_paths)
|
||||
content = content + "\n\nReceived files:\n" + file_list
|
||||
|
||||
if not content:
|
||||
self.channel.logger.warning(
|
||||
"Received empty or unsupported message type: {}",
|
||||
chatbot_msg.message_type,
|
||||
)
|
||||
return AckMessage.STATUS_OK, "OK"
|
||||
|
||||
sender_id = chatbot_msg.sender_staff_id or chatbot_msg.sender_id
|
||||
sender_name = chatbot_msg.sender_nick or "Unknown"
|
||||
|
||||
conversation_type = message.data.get("conversationType")
|
||||
conversation_id = (
|
||||
message.data.get("conversationId")
|
||||
or message.data.get("openConversationId")
|
||||
)
|
||||
|
||||
self.channel.logger.info("Received message from {} ({}): {}", sender_name, sender_id, content)
|
||||
|
||||
# Forward to Vibe-Trading via _on_message (non-blocking).
|
||||
# Store reference to prevent GC before task completes.
|
||||
task = asyncio.create_task(
|
||||
self.channel._on_message(
|
||||
content,
|
||||
sender_id,
|
||||
sender_name,
|
||||
conversation_type,
|
||||
conversation_id,
|
||||
)
|
||||
)
|
||||
self.channel._background_tasks.add(task)
|
||||
task.add_done_callback(self.channel._background_tasks.discard)
|
||||
|
||||
return AckMessage.STATUS_OK, "OK"
|
||||
|
||||
except Exception:
|
||||
self.channel.logger.exception("Error processing message")
|
||||
# Return OK to avoid retry loop from DingTalk server
|
||||
return AckMessage.STATUS_OK, "Error"
|
||||
|
||||
|
||||
class DingTalkConfig(BaseModel):
|
||||
"""DingTalk channel configuration using Stream mode."""
|
||||
|
||||
enabled: bool = False
|
||||
client_id: str = ""
|
||||
client_secret: str = ""
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
allow_remote_media_redirects: bool = False
|
||||
remote_media_redirect_allowed_hosts: list[str] = Field(default_factory=list)
|
||||
group_user_isolation: bool = False # If True, each user in group chat gets their own session
|
||||
|
||||
|
||||
class DingTalkChannel(BaseChannel):
|
||||
"""
|
||||
DingTalk channel using Stream Mode.
|
||||
|
||||
Uses WebSocket to receive events via `dingtalk-stream` SDK.
|
||||
Uses direct HTTP API to send messages (SDK is mainly for receiving).
|
||||
|
||||
Supports both private (1:1) and group chats.
|
||||
Group chat_id is stored with a "group:" prefix to route replies back.
|
||||
"""
|
||||
|
||||
name = "dingtalk"
|
||||
display_name = "DingTalk"
|
||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"}
|
||||
_AUDIO_EXTS = {".amr", ".mp3", ".wav", ".ogg", ".m4a", ".aac"}
|
||||
_VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm"}
|
||||
_ZIP_BEFORE_UPLOAD_EXTS = {".htm", ".html"}
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return DingTalkConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = DingTalkConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: DingTalkConfig = config
|
||||
self._client: Any = None
|
||||
self._http: httpx.AsyncClient | None = None
|
||||
|
||||
# Access Token management for sending messages
|
||||
self._access_token: str | None = None
|
||||
self._token_expiry: float = 0
|
||||
|
||||
# Hold references to background tasks to prevent GC
|
||||
self._background_tasks: set[asyncio.Task] = set()
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the DingTalk bot with Stream Mode."""
|
||||
try:
|
||||
if not DINGTALK_AVAILABLE:
|
||||
self.logger.error(
|
||||
"Stream SDK not installed. Run: pip install dingtalk-stream"
|
||||
)
|
||||
return
|
||||
|
||||
if not self.config.client_id or not self.config.client_secret:
|
||||
self.logger.error("client_id and client_secret not configured")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._http = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(10.0, connect=10.0, read=30.0, write=30.0, pool=10.0)
|
||||
)
|
||||
|
||||
self.logger.info(
|
||||
"Initializing Stream Client with Client ID: {}...",
|
||||
self.config.client_id,
|
||||
)
|
||||
credential = Credential(self.config.client_id, self.config.client_secret)
|
||||
self._client = DingTalkStreamClient(credential)
|
||||
|
||||
# Register standard handler
|
||||
handler = VibeTradingDingTalkHandler(self)
|
||||
self._client.register_callback_handler(ChatbotMessage.TOPIC, handler)
|
||||
|
||||
self.logger.info("bot started with Stream Mode")
|
||||
|
||||
# Reconnect loop: restart stream if SDK exits or crashes
|
||||
while self._running:
|
||||
try:
|
||||
await self._client.start()
|
||||
except Exception as e:
|
||||
self.logger.warning("stream error: {}", e)
|
||||
if self._running:
|
||||
self.logger.info("Reconnecting stream in 5 seconds...")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Failed to start channel")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the DingTalk bot."""
|
||||
self._running = False
|
||||
# Close the shared HTTP client
|
||||
if self._http:
|
||||
await self._http.aclose()
|
||||
self._http = None
|
||||
# Cancel outstanding background tasks
|
||||
for task in self._background_tasks:
|
||||
task.cancel()
|
||||
self._background_tasks.clear()
|
||||
|
||||
async def _get_access_token(self) -> str | None:
|
||||
"""Get or refresh Access Token."""
|
||||
if self._access_token and time.time() < self._token_expiry:
|
||||
return self._access_token
|
||||
|
||||
url = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
|
||||
data = {
|
||||
"appKey": self.config.client_id,
|
||||
"appSecret": self.config.client_secret,
|
||||
}
|
||||
|
||||
if not self._http:
|
||||
self.logger.warning("HTTP client not initialized, cannot refresh token")
|
||||
return None
|
||||
|
||||
try:
|
||||
resp = await self._http.post(url, json=data)
|
||||
resp.raise_for_status()
|
||||
res_data = resp.json()
|
||||
self._access_token = res_data.get("accessToken")
|
||||
# Expire 60s early to be safe
|
||||
self._token_expiry = time.time() + int(res_data.get("expireIn", 7200)) - 60
|
||||
return self._access_token
|
||||
except Exception:
|
||||
self.logger.exception("Failed to get access token")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _is_http_url(value: str) -> bool:
|
||||
return urlparse(value).scheme in ("http", "https")
|
||||
|
||||
def _guess_upload_type(self, media_ref: str) -> str:
|
||||
ext = Path(urlparse(media_ref).path).suffix.lower()
|
||||
if ext in self._IMAGE_EXTS:
|
||||
return "image"
|
||||
if ext in self._AUDIO_EXTS:
|
||||
return "voice"
|
||||
if ext in self._VIDEO_EXTS:
|
||||
return "video"
|
||||
return "file"
|
||||
|
||||
def _guess_filename(self, media_ref: str, upload_type: str) -> str:
|
||||
name = os.path.basename(urlparse(media_ref).path)
|
||||
return name or {"image": "image.jpg", "voice": "audio.amr", "video": "video.mp4"}.get(upload_type, "file.bin")
|
||||
|
||||
@staticmethod
|
||||
def _zip_bytes(filename: str, data: bytes) -> tuple[bytes, str, str]:
|
||||
stem = Path(filename).stem or "attachment"
|
||||
safe_name = filename or "attachment.bin"
|
||||
zip_name = f"{stem}.zip"
|
||||
buffer = BytesIO()
|
||||
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
|
||||
archive.writestr(safe_name, data)
|
||||
return buffer.getvalue(), zip_name, "application/zip"
|
||||
|
||||
def _normalize_upload_payload(
|
||||
self,
|
||||
filename: str,
|
||||
data: bytes,
|
||||
content_type: str | None,
|
||||
) -> tuple[bytes, str, str | None]:
|
||||
ext = Path(filename).suffix.lower()
|
||||
if ext in self._ZIP_BEFORE_UPLOAD_EXTS or content_type == "text/html":
|
||||
self.logger.info(
|
||||
"does not accept raw HTML attachments, zipping {} before upload",
|
||||
filename,
|
||||
)
|
||||
return self._zip_bytes(filename, data)
|
||||
return data, filename, content_type
|
||||
|
||||
def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||
ok, err = validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
|
||||
return False
|
||||
return True
|
||||
|
||||
def _redirect_host_allowed(self, current_url: str, next_url: str) -> bool:
|
||||
current_host = (urlparse(current_url).hostname or "").lower()
|
||||
next_host = (urlparse(next_url).hostname or "").lower()
|
||||
if not next_host:
|
||||
return False
|
||||
if next_host == current_host:
|
||||
return True
|
||||
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
|
||||
return next_host in allowed_hosts
|
||||
|
||||
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None:
|
||||
if not self.config.allow_remote_media_redirects:
|
||||
self.logger.warning("media download redirect refused ref={}", current_url)
|
||||
return None
|
||||
if not location:
|
||||
self.logger.warning("media download redirect without Location ref={}", current_url)
|
||||
return None
|
||||
next_url = urljoin(current_url, location)
|
||||
if not self._redirect_host_allowed(current_url, next_url):
|
||||
self.logger.warning(
|
||||
"media download cross-host redirect refused ref={} next={}",
|
||||
current_url,
|
||||
next_url,
|
||||
)
|
||||
return None
|
||||
if not self._validate_remote_media_url(next_url):
|
||||
return None
|
||||
return next_url
|
||||
|
||||
async def _fetch_remote_media_bytes(
|
||||
self,
|
||||
media_ref: str,
|
||||
) -> tuple[bytes | None, str | None]:
|
||||
"""Fetch a remote media URL with SSRF, redirect, and size checks."""
|
||||
if not self._http:
|
||||
return None, None
|
||||
|
||||
if not self._validate_remote_media_url(media_ref):
|
||||
return None, None
|
||||
|
||||
try:
|
||||
# Prefer streaming with a running byte cap so large responses are not
|
||||
# materialized before the limit is enforced. Test fakes may only
|
||||
# implement get(), so keep a small compatibility fallback below.
|
||||
stream = getattr(self._http, "stream", None)
|
||||
if stream is not None:
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
async with stream("GET", current_url, follow_redirects=False) as resp:
|
||||
final_ok, final_err = validate_resolved_url(str(resp.url))
|
||||
if not final_ok:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
media_ref,
|
||||
resp.url,
|
||||
final_err,
|
||||
)
|
||||
return None, None
|
||||
if 300 <= resp.status_code < 400:
|
||||
next_url = self._next_remote_media_url(
|
||||
str(resp.url), resp.headers.get("location")
|
||||
)
|
||||
if not next_url:
|
||||
return None, None
|
||||
current_url = next_url
|
||||
continue
|
||||
if resp.status_code >= 400:
|
||||
self.logger.warning(
|
||||
"media download failed status={} ref={}",
|
||||
resp.status_code,
|
||||
current_url,
|
||||
)
|
||||
return None, None
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
async for chunk in resp.aiter_bytes():
|
||||
total += len(chunk)
|
||||
if total > DINGTALK_MAX_REMOTE_MEDIA_BYTES:
|
||||
self.logger.warning(
|
||||
"media download too large ref={} bytes>{}",
|
||||
current_url,
|
||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES,
|
||||
)
|
||||
return None, None
|
||||
chunks.append(chunk)
|
||||
return b"".join(chunks), (resp.headers.get("content-type") or "")
|
||||
self.logger.warning("media download exceeded redirect limit ref={}", media_ref)
|
||||
return None, None
|
||||
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
resp = await self._http.get(current_url, follow_redirects=False)
|
||||
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url)))
|
||||
if not final_ok:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
media_ref,
|
||||
getattr(resp, "url", current_url),
|
||||
final_err,
|
||||
)
|
||||
return None, None
|
||||
if 300 <= resp.status_code < 400:
|
||||
next_url = self._next_remote_media_url(
|
||||
str(getattr(resp, "url", current_url)), resp.headers.get("location")
|
||||
)
|
||||
if not next_url:
|
||||
return None, None
|
||||
current_url = next_url
|
||||
continue
|
||||
if resp.status_code >= 400:
|
||||
self.logger.warning(
|
||||
"media download failed status={} ref={}",
|
||||
resp.status_code,
|
||||
current_url,
|
||||
)
|
||||
return None, None
|
||||
if len(resp.content) > DINGTALK_MAX_REMOTE_MEDIA_BYTES:
|
||||
self.logger.warning(
|
||||
"media download too large ref={} bytes>{}",
|
||||
current_url,
|
||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES,
|
||||
)
|
||||
return None, None
|
||||
return resp.content, (resp.headers.get("content-type") or "")
|
||||
self.logger.warning("media download exceeded redirect limit ref={}", media_ref)
|
||||
return None, None
|
||||
except httpx.TransportError:
|
||||
self.logger.exception("media download network error ref={}", media_ref)
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.exception("media download error ref={}", media_ref)
|
||||
return None, None
|
||||
|
||||
async def _read_media_bytes(
|
||||
self,
|
||||
media_ref: str,
|
||||
) -> tuple[bytes | None, str | None, str | None]:
|
||||
if not media_ref:
|
||||
return None, None, None
|
||||
|
||||
if self._is_http_url(media_ref):
|
||||
data, raw_content_type = await self._fetch_remote_media_bytes(media_ref)
|
||||
if data is None:
|
||||
return None, None, None
|
||||
content_type = (raw_content_type or "").split(";")[0].strip()
|
||||
filename = self._guess_filename(media_ref, self._guess_upload_type(media_ref))
|
||||
return data, filename, content_type or None
|
||||
|
||||
try:
|
||||
if media_ref.startswith("file://"):
|
||||
parsed = urlparse(media_ref)
|
||||
local_path = Path(unquote(parsed.path))
|
||||
else:
|
||||
local_path = Path(os.path.expanduser(media_ref))
|
||||
if not local_path.is_file():
|
||||
self.logger.warning("media file not found: {}", local_path)
|
||||
return None, None, None
|
||||
data = await asyncio.to_thread(local_path.read_bytes)
|
||||
content_type = mimetypes.guess_type(local_path.name)[0]
|
||||
return data, local_path.name, content_type
|
||||
except Exception:
|
||||
self.logger.exception("media read error ref={}", media_ref)
|
||||
return None, None, None
|
||||
|
||||
async def _upload_media(
|
||||
self,
|
||||
token: str,
|
||||
data: bytes,
|
||||
media_type: str,
|
||||
filename: str,
|
||||
content_type: str | None,
|
||||
) -> str | None:
|
||||
if not self._http:
|
||||
return None
|
||||
url = f"https://oapi.dingtalk.com/media/upload?access_token={token}&type={media_type}"
|
||||
mime = content_type or mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||||
files = {"media": (filename, data, mime)}
|
||||
|
||||
try:
|
||||
resp = await self._http.post(url, files=files)
|
||||
text = resp.text
|
||||
result = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
|
||||
if resp.status_code >= 400:
|
||||
self.logger.error("media upload failed status={} type={} body={}", resp.status_code, media_type, text[:500])
|
||||
return None
|
||||
errcode = result.get("errcode", 0)
|
||||
if errcode != 0:
|
||||
self.logger.error("media upload api error type={} errcode={} body={}", media_type, errcode, text[:500])
|
||||
return None
|
||||
sub = result.get("result") or {}
|
||||
media_id = result.get("media_id") or result.get("mediaId") or sub.get("media_id") or sub.get("mediaId")
|
||||
if not media_id:
|
||||
self.logger.error("media upload missing media_id body={}", text[:500])
|
||||
return None
|
||||
return str(media_id)
|
||||
except httpx.TransportError:
|
||||
self.logger.exception("media upload network error type={}", media_type)
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.exception("media upload error type={}", media_type)
|
||||
return None
|
||||
|
||||
async def _send_batch_message(
|
||||
self,
|
||||
token: str,
|
||||
chat_id: str,
|
||||
msg_key: str,
|
||||
msg_param: dict[str, Any],
|
||||
) -> bool:
|
||||
if not self._http:
|
||||
self.logger.warning("HTTP client not initialized, cannot send")
|
||||
return False
|
||||
|
||||
headers = {"x-acs-dingtalk-access-token": token}
|
||||
if chat_id.startswith("group:"):
|
||||
# Group chat
|
||||
url = "https://api.dingtalk.com/v1.0/robot/groupMessages/send"
|
||||
payload = {
|
||||
"robotCode": self.config.client_id,
|
||||
"openConversationId": chat_id[6:], # Remove "group:" prefix,
|
||||
"msgKey": msg_key,
|
||||
"msgParam": json.dumps(msg_param, ensure_ascii=False),
|
||||
}
|
||||
else:
|
||||
# Private chat
|
||||
url = "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend"
|
||||
payload = {
|
||||
"robotCode": self.config.client_id,
|
||||
"userIds": [chat_id],
|
||||
"msgKey": msg_key,
|
||||
"msgParam": json.dumps(msg_param, ensure_ascii=False),
|
||||
}
|
||||
|
||||
try:
|
||||
resp = await self._http.post(url, json=payload, headers=headers)
|
||||
body = resp.text
|
||||
if resp.status_code != 200:
|
||||
self.logger.error("send failed msgKey={} status={} body={}", msg_key, resp.status_code, body[:500])
|
||||
return False
|
||||
try:
|
||||
result = resp.json()
|
||||
except Exception:
|
||||
result = {}
|
||||
errcode = result.get("errcode")
|
||||
if errcode not in (None, 0):
|
||||
self.logger.error("send api error msgKey={} errcode={} body={}", msg_key, errcode, body[:500])
|
||||
return False
|
||||
self.logger.debug("message sent to {} with msgKey={}", chat_id, msg_key)
|
||||
return True
|
||||
except httpx.TransportError:
|
||||
self.logger.exception("network error sending message msgKey={}", msg_key)
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.exception("Error sending message msgKey={}", msg_key)
|
||||
return False
|
||||
|
||||
async def _send_markdown_text(self, token: str, chat_id: str, content: str) -> bool:
|
||||
return await self._send_batch_message(
|
||||
token,
|
||||
chat_id,
|
||||
"sampleMarkdown",
|
||||
{"text": content, "title": "Vibe-Trading Reply"},
|
||||
)
|
||||
|
||||
async def _send_media_ref(self, token: str, chat_id: str, media_ref: str) -> bool:
|
||||
media_ref = (media_ref or "").strip()
|
||||
if not media_ref:
|
||||
return True
|
||||
|
||||
upload_type = self._guess_upload_type(media_ref)
|
||||
if upload_type == "image" and self._is_http_url(media_ref):
|
||||
ok = await self._send_batch_message(
|
||||
token,
|
||||
chat_id,
|
||||
"sampleImageMsg",
|
||||
{"photoURL": media_ref},
|
||||
)
|
||||
if ok:
|
||||
return True
|
||||
self.logger.warning("image url send failed, trying upload fallback: {}", media_ref)
|
||||
|
||||
data, filename, content_type = await self._read_media_bytes(media_ref)
|
||||
if not data:
|
||||
self.logger.error("media read failed: {}", media_ref)
|
||||
return False
|
||||
|
||||
filename = filename or self._guess_filename(media_ref, upload_type)
|
||||
data, filename, content_type = self._normalize_upload_payload(filename, data, content_type)
|
||||
file_type = Path(filename).suffix.lower().lstrip(".")
|
||||
if not file_type:
|
||||
guessed = mimetypes.guess_extension(content_type or "")
|
||||
file_type = (guessed or ".bin").lstrip(".")
|
||||
if file_type == "jpeg":
|
||||
file_type = "jpg"
|
||||
|
||||
media_id = await self._upload_media(
|
||||
token=token,
|
||||
data=data,
|
||||
media_type=upload_type,
|
||||
filename=filename,
|
||||
content_type=content_type,
|
||||
)
|
||||
if not media_id:
|
||||
return False
|
||||
|
||||
if upload_type == "image":
|
||||
# Verified in production: sampleImageMsg accepts media_id in photoURL.
|
||||
ok = await self._send_batch_message(
|
||||
token,
|
||||
chat_id,
|
||||
"sampleImageMsg",
|
||||
{"photoURL": media_id},
|
||||
)
|
||||
if ok:
|
||||
return True
|
||||
self.logger.warning("image media_id send failed, falling back to file: {}", media_ref)
|
||||
|
||||
return await self._send_batch_message(
|
||||
token,
|
||||
chat_id,
|
||||
"sampleFile",
|
||||
{"mediaId": media_id, "fileName": filename, "fileType": file_type},
|
||||
)
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a message through DingTalk."""
|
||||
token = await self._get_access_token()
|
||||
if not token:
|
||||
return
|
||||
|
||||
if msg.content and msg.content.strip():
|
||||
await self._send_markdown_text(token, msg.chat_id, msg.content.strip())
|
||||
|
||||
for media_ref in msg.media or []:
|
||||
ok = await self._send_media_ref(token, msg.chat_id, media_ref)
|
||||
if ok:
|
||||
continue
|
||||
self.logger.error("media send failed for {}", media_ref)
|
||||
# Send visible fallback so failures are observable by the user.
|
||||
filename = self._guess_filename(media_ref, self._guess_upload_type(media_ref))
|
||||
await self._send_markdown_text(
|
||||
token,
|
||||
msg.chat_id,
|
||||
f"[Attachment send failed: {filename}]",
|
||||
)
|
||||
|
||||
async def _on_message(
|
||||
self,
|
||||
content: str,
|
||||
sender_id: str,
|
||||
sender_name: str,
|
||||
conversation_type: str | None = None,
|
||||
conversation_id: str | None = None,
|
||||
) -> None:
|
||||
"""Handle incoming message (called by VibeTradingDingTalkHandler).
|
||||
|
||||
Delegates to BaseChannel._handle_message() which enforces allow_from
|
||||
permission checks before publishing to the bus.
|
||||
"""
|
||||
try:
|
||||
self.logger.info("inbound: {} from {}", content, sender_name)
|
||||
is_group = conversation_type == "2" and conversation_id
|
||||
chat_id = f"group:{conversation_id}" if is_group else sender_id
|
||||
session_key = None
|
||||
if is_group and self.config.group_user_isolation:
|
||||
session_key = f"{self.name}:group:{conversation_id}:{sender_id}"
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
content=str(content),
|
||||
metadata={
|
||||
"sender_name": sender_name,
|
||||
"platform": "dingtalk",
|
||||
"conversation_type": conversation_type,
|
||||
},
|
||||
session_key=session_key,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Error publishing message")
|
||||
|
||||
async def _download_dingtalk_file(
|
||||
self,
|
||||
download_code: str,
|
||||
filename: str,
|
||||
sender_id: str,
|
||||
) -> str | None:
|
||||
"""Download a DingTalk file to the media directory, return local path."""
|
||||
from src.channels.utils import get_media_dir
|
||||
|
||||
try:
|
||||
token = await self._get_access_token()
|
||||
if not token or not self._http:
|
||||
self.logger.error("file download: no token or http client")
|
||||
return None
|
||||
|
||||
# Step 1: Exchange downloadCode for a temporary download URL
|
||||
api_url = "https://api.dingtalk.com/v1.0/robot/messageFiles/download"
|
||||
headers = {"x-acs-dingtalk-access-token": token, "Content-Type": "application/json"}
|
||||
payload = {"downloadCode": download_code, "robotCode": self.config.client_id}
|
||||
resp = await self._http.post(api_url, json=payload, headers=headers)
|
||||
if resp.status_code != 200:
|
||||
self.logger.error("get download URL failed: status={}, body={}", resp.status_code, resp.text)
|
||||
return None
|
||||
|
||||
result = resp.json()
|
||||
download_url = result.get("downloadUrl")
|
||||
if not download_url:
|
||||
self.logger.error("download URL not found in response: {}", result)
|
||||
return None
|
||||
|
||||
# Step 2: Download the file content
|
||||
file_resp = await self._http.get(download_url, follow_redirects=True)
|
||||
if file_resp.status_code != 200:
|
||||
self.logger.error("file download failed: status={}", file_resp.status_code)
|
||||
return None
|
||||
|
||||
# Save to media directory (accessible under workspace)
|
||||
download_dir = get_media_dir("dingtalk") / sender_id
|
||||
download_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = download_dir / filename
|
||||
await asyncio.to_thread(file_path.write_bytes, file_resp.content)
|
||||
self.logger.info("file saved: {}", file_path)
|
||||
return str(file_path)
|
||||
except Exception:
|
||||
self.logger.exception("file download error")
|
||||
return None
|
||||
@@ -0,0 +1,818 @@
|
||||
"""Discord channel implementation using discord.py."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
|
||||
from src.channels.utils import get_media_dir
|
||||
from pydantic import BaseModel
|
||||
from src.channels.utils import safe_filename, split_message
|
||||
|
||||
DISCORD_AVAILABLE = importlib.util.find_spec("discord") is not None
|
||||
if TYPE_CHECKING:
|
||||
import aiohttp
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.abc import Messageable
|
||||
|
||||
if DISCORD_AVAILABLE:
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.abc import Messageable
|
||||
|
||||
MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024 # 20MB
|
||||
MAX_MESSAGE_LEN = 2000 # Discord message character limit
|
||||
TYPING_INTERVAL_S = 8
|
||||
|
||||
|
||||
@dataclass
|
||||
class _StreamBuf:
|
||||
"""Per-chat streaming accumulator for progressive Discord message edits."""
|
||||
|
||||
text: str = ""
|
||||
message: Any | None = None
|
||||
last_edit: float = 0.0
|
||||
stream_id: str | None = None
|
||||
|
||||
|
||||
class DiscordConfig(BaseModel):
|
||||
"""Discord channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
token: str = ""
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
allow_channels: list[str] = Field(default_factory=list) # Allowed channel IDs (empty = all)
|
||||
intents: int = 37377
|
||||
group_policy: Literal["mention", "open"] = "mention"
|
||||
read_receipt_emoji: str = "👀"
|
||||
working_emoji: str = "🔧"
|
||||
working_emoji_delay: float = 2.0
|
||||
streaming: bool = True
|
||||
proxy: str | None = None
|
||||
proxy_username: str | None = None
|
||||
proxy_password: str | None = None
|
||||
|
||||
|
||||
if DISCORD_AVAILABLE:
|
||||
|
||||
class DiscordBotClient(discord.Client):
|
||||
"""discord.py client that forwards events to the channel."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channel: DiscordChannel,
|
||||
*,
|
||||
intents: discord.Intents,
|
||||
proxy: str | None = None,
|
||||
proxy_auth: aiohttp.BasicAuth | None = None,
|
||||
) -> None:
|
||||
super().__init__(intents=intents, proxy=proxy, proxy_auth=proxy_auth)
|
||||
self._channel = channel
|
||||
self.tree = app_commands.CommandTree(self)
|
||||
self._register_app_commands()
|
||||
|
||||
async def on_ready(self) -> None:
|
||||
self._channel._bot_user_id = str(self.user.id) if self.user else None
|
||||
self._channel.logger.info("bot connected as user {}", self._channel._bot_user_id)
|
||||
try:
|
||||
synced = await self.tree.sync()
|
||||
self._channel.logger.info("app commands synced: {}", len(synced))
|
||||
except Exception as e:
|
||||
self._channel.logger.warning("app command sync failed: {}", e)
|
||||
|
||||
async def on_message(self, message: discord.Message) -> None:
|
||||
await self._channel._handle_discord_message(message)
|
||||
|
||||
async def on_thread_delete(self, thread: discord.Thread) -> None:
|
||||
self._channel._forget_channel(thread)
|
||||
|
||||
async def on_thread_update(self, before: discord.Thread, after: discord.Thread) -> None:
|
||||
if getattr(after, "archived", False):
|
||||
self._channel._forget_channel(after)
|
||||
else:
|
||||
self._channel._remember_channel(after)
|
||||
|
||||
async def _reply_ephemeral(self, interaction: discord.Interaction, text: str) -> bool:
|
||||
"""Send an ephemeral interaction response and report success."""
|
||||
try:
|
||||
await interaction.response.send_message(text, ephemeral=True)
|
||||
return True
|
||||
except Exception as e:
|
||||
self._channel.logger.warning("interaction response failed: {}", e)
|
||||
return False
|
||||
|
||||
async def _resolve_interaction_channel(
|
||||
self,
|
||||
interaction: discord.Interaction,
|
||||
) -> Any | None:
|
||||
channel_id = interaction.channel_id
|
||||
if channel_id is None:
|
||||
return None
|
||||
channel = getattr(interaction, "channel", None) or self.get_channel(channel_id)
|
||||
if channel is None:
|
||||
try:
|
||||
channel = await self.fetch_channel(channel_id)
|
||||
except Exception as e:
|
||||
self._channel.logger.warning("interaction channel {} unavailable: {}", channel_id, e)
|
||||
return None
|
||||
self._channel._remember_channel(channel)
|
||||
return channel
|
||||
|
||||
async def _interaction_channel_allowed(
|
||||
self,
|
||||
interaction: discord.Interaction,
|
||||
channel: Any | None,
|
||||
) -> bool:
|
||||
allow_channels = self._channel.config.allow_channels
|
||||
if not allow_channels:
|
||||
return True
|
||||
if channel is None:
|
||||
channel_id = interaction.channel_id
|
||||
return channel_id is not None and str(channel_id) in allow_channels
|
||||
channel_ids = self._channel._channel_allow_keys(channel)
|
||||
return not channel_ids.isdisjoint(allow_channels)
|
||||
|
||||
async def _forward_slash_command(
|
||||
self,
|
||||
interaction: discord.Interaction,
|
||||
command_text: str,
|
||||
) -> None:
|
||||
sender_id = str(interaction.user.id)
|
||||
channel_id = interaction.channel_id
|
||||
|
||||
if channel_id is None:
|
||||
self._channel.logger.warning("slash command missing channel_id: {}", command_text)
|
||||
return
|
||||
|
||||
if not self._channel.is_allowed(sender_id):
|
||||
await self._reply_ephemeral(interaction, "You are not allowed to use this bot.")
|
||||
return
|
||||
|
||||
channel = await self._resolve_interaction_channel(interaction)
|
||||
if not await self._interaction_channel_allowed(interaction, channel):
|
||||
await self._reply_ephemeral(interaction, "This channel is not allowed for this bot.")
|
||||
return
|
||||
|
||||
await self._reply_ephemeral(interaction, f"Processing {command_text}...")
|
||||
|
||||
metadata: dict[str, Any] = {
|
||||
"interaction_id": str(interaction.id),
|
||||
"guild_id": str(interaction.guild_id) if interaction.guild_id else None,
|
||||
"is_slash_command": True,
|
||||
}
|
||||
session_key = None
|
||||
if channel is not None:
|
||||
parent_channel_id = self._channel._channel_parent_key(channel)
|
||||
if parent_channel_id is not None:
|
||||
metadata["parent_channel_id"] = parent_channel_id
|
||||
metadata["context_chat_id"] = parent_channel_id
|
||||
metadata["thread_id"] = str(channel_id)
|
||||
session_key = f"{self._channel.name}:{parent_channel_id}:thread:{channel_id}"
|
||||
|
||||
await self._channel._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=str(channel_id),
|
||||
content=command_text,
|
||||
metadata=metadata,
|
||||
session_key=session_key,
|
||||
)
|
||||
|
||||
def _register_app_commands(self) -> None:
|
||||
commands = (
|
||||
("new", "Stop current task and start a new conversation", "/new"),
|
||||
("stop", "Stop the current task", "/stop"),
|
||||
("restart", "Restart the bot", "/restart"),
|
||||
("status", "Show bot status", "/status"),
|
||||
("history", "Show recent conversation messages", "/history"),
|
||||
)
|
||||
|
||||
for name, description, command_text in commands:
|
||||
|
||||
@self.tree.command(name=name, description=description)
|
||||
async def command_handler(
|
||||
interaction: discord.Interaction,
|
||||
_command_text: str = command_text,
|
||||
) -> None:
|
||||
await self._forward_slash_command(interaction, _command_text)
|
||||
|
||||
@self.tree.command(name="model", description="Show or switch runtime model preset")
|
||||
@app_commands.describe(preset="Optional model preset name, such as default")
|
||||
async def model_command(
|
||||
interaction: discord.Interaction,
|
||||
preset: str | None = None,
|
||||
) -> None:
|
||||
preset = (preset or "").strip()
|
||||
command_text = f"/model {preset}" if preset else "/model"
|
||||
await self._forward_slash_command(interaction, command_text)
|
||||
|
||||
@self.tree.command(name="help", description="Show available commands")
|
||||
async def help_command(interaction: discord.Interaction) -> None:
|
||||
sender_id = str(interaction.user.id)
|
||||
if not self._channel.is_allowed(sender_id):
|
||||
await self._reply_ephemeral(interaction, "You are not allowed to use this bot.")
|
||||
return
|
||||
channel = await self._resolve_interaction_channel(interaction)
|
||||
if not await self._interaction_channel_allowed(interaction, channel):
|
||||
await self._reply_ephemeral(interaction, "This channel is not allowed for this bot.")
|
||||
return
|
||||
await self._reply_ephemeral(interaction, "Use /help for available commands. Docs: https://github.com/HKUDS/Vibe-Trading")
|
||||
|
||||
@self.tree.error
|
||||
async def on_app_command_error(
|
||||
interaction: discord.Interaction,
|
||||
error: app_commands.AppCommandError,
|
||||
) -> None:
|
||||
command_name = interaction.command.qualified_name if interaction.command else "?"
|
||||
self._channel.logger.warning(
|
||||
"app command failed user={} channel={} cmd={} error={}",
|
||||
interaction.user.id,
|
||||
interaction.channel_id,
|
||||
command_name,
|
||||
error,
|
||||
)
|
||||
|
||||
async def send_outbound(self, msg: OutboundMessage) -> None:
|
||||
"""Send a vibe-trading outbound message using Discord transport rules."""
|
||||
channel_id = int(msg.chat_id)
|
||||
|
||||
channel = self._channel._known_channels.get(msg.chat_id) or self.get_channel(channel_id)
|
||||
if channel is None:
|
||||
try:
|
||||
channel = await self.fetch_channel(channel_id)
|
||||
except Exception as e:
|
||||
self._channel.logger.warning("channel {} unavailable: {}", msg.chat_id, e)
|
||||
return
|
||||
|
||||
reference, mention_settings = self._build_reply_context(channel, msg.reply_to)
|
||||
sent_media = False
|
||||
failed_media: list[str] = []
|
||||
|
||||
for index, media_path in enumerate(msg.media or []):
|
||||
if await self._send_file(
|
||||
channel,
|
||||
media_path,
|
||||
reference=reference if index == 0 else None,
|
||||
mention_settings=mention_settings,
|
||||
):
|
||||
sent_media = True
|
||||
else:
|
||||
failed_media.append(Path(media_path).name)
|
||||
|
||||
for index, chunk in enumerate(
|
||||
self._build_chunks(msg.content or "", failed_media, sent_media)
|
||||
):
|
||||
kwargs: dict[str, Any] = {"content": chunk}
|
||||
if index == 0 and reference is not None and not sent_media:
|
||||
kwargs["reference"] = reference
|
||||
kwargs["allowed_mentions"] = mention_settings
|
||||
await channel.send(**kwargs)
|
||||
|
||||
async def _send_file(
|
||||
self,
|
||||
channel: Messageable,
|
||||
file_path: str,
|
||||
*,
|
||||
reference: discord.PartialMessage | None,
|
||||
mention_settings: discord.AllowedMentions,
|
||||
) -> bool:
|
||||
"""Send a file attachment via discord.py."""
|
||||
path = Path(file_path)
|
||||
if not path.is_file():
|
||||
self._channel.logger.warning("file not found, skipping: {}", file_path)
|
||||
return False
|
||||
|
||||
if path.stat().st_size > MAX_ATTACHMENT_BYTES:
|
||||
self._channel.logger.warning("file too large (>20MB), skipping: {}", path.name)
|
||||
return False
|
||||
|
||||
try:
|
||||
kwargs: dict[str, Any] = {"file": discord.File(path)}
|
||||
if reference is not None:
|
||||
kwargs["reference"] = reference
|
||||
kwargs["allowed_mentions"] = mention_settings
|
||||
await channel.send(**kwargs)
|
||||
self._channel.logger.info("file sent: {}", path.name)
|
||||
return True
|
||||
except Exception:
|
||||
self._channel.logger.exception("Error sending file {}", path.name)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _build_chunks(content: str, failed_media: list[str], sent_media: bool) -> list[str]:
|
||||
"""Build outbound text chunks, including attachment-failure fallback text."""
|
||||
chunks = split_message(content, MAX_MESSAGE_LEN)
|
||||
if chunks or not failed_media or sent_media:
|
||||
return chunks
|
||||
fallback = "\n".join(f"[attachment: {name} - send failed]" for name in failed_media)
|
||||
return split_message(fallback, MAX_MESSAGE_LEN)
|
||||
|
||||
def _build_reply_context(
|
||||
self,
|
||||
channel: Messageable,
|
||||
reply_to: str | None,
|
||||
) -> tuple[discord.PartialMessage | None, discord.AllowedMentions]:
|
||||
"""Build reply context for outbound messages."""
|
||||
mention_settings = discord.AllowedMentions(replied_user=False)
|
||||
if not reply_to:
|
||||
return None, mention_settings
|
||||
try:
|
||||
message_id = int(reply_to)
|
||||
except ValueError:
|
||||
self._channel.logger.warning("Invalid reply target: {}", reply_to)
|
||||
return None, mention_settings
|
||||
|
||||
return channel.get_partial_message(message_id), mention_settings
|
||||
|
||||
|
||||
class DiscordChannel(BaseChannel):
|
||||
"""Discord channel using discord.py."""
|
||||
|
||||
name = "discord"
|
||||
display_name = "Discord"
|
||||
_STREAM_EDIT_INTERVAL = 0.8
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return DiscordConfig().model_dump(by_alias=True)
|
||||
|
||||
@staticmethod
|
||||
def _channel_key(channel_or_id: Any) -> str:
|
||||
"""Normalize channel-like objects and ids to a stable string key."""
|
||||
channel_id = getattr(channel_or_id, "id", channel_or_id)
|
||||
return str(channel_id)
|
||||
|
||||
@classmethod
|
||||
def _channel_allow_keys(cls, channel: Any) -> set[str]:
|
||||
"""Return channel IDs that can satisfy allow_channels for this channel."""
|
||||
keys = {cls._channel_key(channel)}
|
||||
if parent_key := cls._channel_parent_key(channel):
|
||||
keys.add(parent_key)
|
||||
return keys
|
||||
|
||||
@classmethod
|
||||
def _channel_parent_key(cls, channel: Any) -> str | None:
|
||||
"""Return the parent channel key for a Discord thread-like channel."""
|
||||
parent_id = getattr(channel, "parent_id", None)
|
||||
if parent_id is not None:
|
||||
return cls._channel_key(parent_id)
|
||||
parent = getattr(channel, "parent", None)
|
||||
if parent is not None:
|
||||
return cls._channel_key(parent)
|
||||
return None
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = DiscordConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: DiscordConfig = config
|
||||
self._client: DiscordBotClient | None = None
|
||||
self._typing_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._bot_user_id: str | None = None
|
||||
self._pending_reactions: dict[str, Any] = {} # chat_id -> message object
|
||||
self._working_emoji_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._stream_bufs: dict[str, _StreamBuf] = {}
|
||||
self._known_channels: dict[str, Any] = {}
|
||||
|
||||
def _remember_channel(self, channel: Any) -> None:
|
||||
self._known_channels[self._channel_key(channel)] = channel
|
||||
|
||||
def _forget_channel(self, channel_or_id: Any) -> None:
|
||||
self._known_channels.pop(self._channel_key(channel_or_id), None)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the Discord client."""
|
||||
if not DISCORD_AVAILABLE:
|
||||
self.logger.error("discord.py not installed. Run: pip install vibe-trading-ai[discord]")
|
||||
return
|
||||
|
||||
if not self.config.token:
|
||||
self.logger.error("bot token not configured")
|
||||
return
|
||||
|
||||
try:
|
||||
intents = discord.Intents.none()
|
||||
intents.value = self.config.intents
|
||||
|
||||
proxy_auth = None
|
||||
has_user = bool(self.config.proxy_username)
|
||||
has_pass = bool(self.config.proxy_password)
|
||||
if has_user and has_pass:
|
||||
import aiohttp
|
||||
|
||||
proxy_auth = aiohttp.BasicAuth(
|
||||
login=self.config.proxy_username,
|
||||
password=self.config.proxy_password,
|
||||
)
|
||||
elif has_user != has_pass:
|
||||
self.logger.warning(
|
||||
"proxy auth incomplete: both proxy_username and "
|
||||
"proxy_password must be set; ignoring partial credentials",
|
||||
)
|
||||
|
||||
self._client = DiscordBotClient(
|
||||
self,
|
||||
intents=intents,
|
||||
proxy=self.config.proxy,
|
||||
proxy_auth=proxy_auth,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to initialize client")
|
||||
self._client = None
|
||||
self._running = False
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self.logger.info("Starting client via discord.py...")
|
||||
|
||||
try:
|
||||
await self._client.start(self.config.token)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.exception("client startup failed")
|
||||
finally:
|
||||
self._running = False
|
||||
await self._reset_runtime_state(close_client=True)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the Discord channel."""
|
||||
self._running = False
|
||||
await self._reset_runtime_state(close_client=True)
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a message through Discord using discord.py."""
|
||||
client = self._client
|
||||
if client is None or not client.is_ready():
|
||||
self.logger.warning("client not ready; dropping outbound message")
|
||||
return
|
||||
|
||||
is_progress = bool((msg.metadata or {}).get("_progress"))
|
||||
|
||||
try:
|
||||
await client.send_outbound(msg)
|
||||
except Exception:
|
||||
self.logger.exception("Error sending message")
|
||||
raise
|
||||
finally:
|
||||
if not is_progress:
|
||||
await self._stop_typing(msg.chat_id)
|
||||
await self._clear_reactions(msg.chat_id)
|
||||
|
||||
async def send_delta(
|
||||
self, chat_id: str, delta: str, metadata: dict[str, Any] | None = None
|
||||
) -> None:
|
||||
"""Progressive Discord delivery: send once, then edit until the stream ends."""
|
||||
client = self._client
|
||||
if client is None or not client.is_ready():
|
||||
self.logger.warning("client not ready; dropping stream delta")
|
||||
return
|
||||
|
||||
meta = metadata or {}
|
||||
stream_id = meta.get("_stream_id")
|
||||
|
||||
if meta.get("_stream_end"):
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if not buf or buf.message is None or not buf.text:
|
||||
return
|
||||
if stream_id is not None and buf.stream_id is not None and buf.stream_id != stream_id:
|
||||
return
|
||||
await self._finalize_stream(chat_id, buf)
|
||||
return
|
||||
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if buf is None or (
|
||||
stream_id is not None and buf.stream_id is not None and buf.stream_id != stream_id
|
||||
):
|
||||
buf = _StreamBuf(stream_id=stream_id)
|
||||
self._stream_bufs[chat_id] = buf
|
||||
elif buf.stream_id is None:
|
||||
buf.stream_id = stream_id
|
||||
|
||||
buf.text += delta
|
||||
if not buf.text.strip():
|
||||
return
|
||||
|
||||
target = await self._resolve_channel(chat_id)
|
||||
if target is None:
|
||||
self.logger.warning("stream target {} unavailable", chat_id)
|
||||
return
|
||||
|
||||
now = time.monotonic()
|
||||
if buf.message is None:
|
||||
try:
|
||||
buf.message = await target.send(content=buf.text)
|
||||
buf.last_edit = now
|
||||
except Exception as e:
|
||||
self.logger.warning("stream initial send failed: {}", e)
|
||||
raise
|
||||
return
|
||||
|
||||
if (now - buf.last_edit) < self._STREAM_EDIT_INTERVAL:
|
||||
return
|
||||
|
||||
try:
|
||||
await buf.message.edit(content=DiscordBotClient._build_chunks(buf.text, [], False)[0])
|
||||
buf.last_edit = now
|
||||
except Exception as e:
|
||||
self.logger.warning("stream edit failed: {}", e)
|
||||
raise
|
||||
|
||||
async def _handle_discord_message(self, message: discord.Message) -> None:
|
||||
"""Handle incoming Discord messages from discord.py.
|
||||
|
||||
Self-loop guard: only drop messages from this bot's own account. Messages
|
||||
from other bots are allowed through so multi-agent setups (one bot asking
|
||||
another for help, a bot mentioning another by @name, etc.) can work.
|
||||
Bot-from-bot loops are still prevented per-instance because each bot
|
||||
still ignores its own outbound messages. (#3217)
|
||||
"""
|
||||
if self._bot_user_id is not None and str(message.author.id) == self._bot_user_id:
|
||||
return
|
||||
if self._is_system_message(message):
|
||||
return
|
||||
|
||||
sender_id = str(message.author.id)
|
||||
channel_id = self._channel_key(message.channel)
|
||||
self._remember_channel(message.channel)
|
||||
content = message.content or ""
|
||||
|
||||
if not self._should_accept_inbound(message, sender_id, content):
|
||||
return
|
||||
|
||||
media_paths, attachment_markers = await self._download_attachments(message.attachments)
|
||||
full_content = self._compose_inbound_content(content, attachment_markers)
|
||||
metadata = self._build_inbound_metadata(message)
|
||||
parent_channel_id = self._channel_parent_key(message.channel)
|
||||
session_key = None
|
||||
if parent_channel_id is not None:
|
||||
metadata["parent_channel_id"] = parent_channel_id
|
||||
metadata["context_chat_id"] = parent_channel_id
|
||||
metadata["thread_id"] = channel_id
|
||||
session_key = f"{self.name}:{parent_channel_id}:thread:{channel_id}"
|
||||
|
||||
await self._start_typing(message.channel)
|
||||
|
||||
# Add read receipt reaction immediately, working emoji after delay
|
||||
try:
|
||||
await message.add_reaction(self.config.read_receipt_emoji)
|
||||
self._pending_reactions[channel_id] = message
|
||||
except Exception as e:
|
||||
self.logger.debug("Failed to add read receipt reaction: {}", e)
|
||||
|
||||
# Delayed working indicator (cosmetic — not tied to subagent lifecycle)
|
||||
async def _delayed_working_emoji() -> None:
|
||||
await asyncio.sleep(self.config.working_emoji_delay)
|
||||
with suppress(Exception):
|
||||
await message.add_reaction(self.config.working_emoji)
|
||||
|
||||
self._working_emoji_tasks[channel_id] = asyncio.create_task(_delayed_working_emoji())
|
||||
|
||||
try:
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=channel_id,
|
||||
content=full_content,
|
||||
media=media_paths,
|
||||
metadata=metadata,
|
||||
session_key=session_key,
|
||||
is_dm=message.guild is None,
|
||||
)
|
||||
except Exception:
|
||||
await self._clear_reactions(channel_id)
|
||||
await self._stop_typing(channel_id)
|
||||
raise
|
||||
|
||||
async def _on_message(self, message: discord.Message) -> None:
|
||||
"""Backward-compatible alias for legacy tests/callers."""
|
||||
await self._handle_discord_message(message)
|
||||
|
||||
async def _resolve_channel(self, chat_id: str) -> Any | None:
|
||||
"""Resolve a Discord channel from cache first, then network fetch."""
|
||||
client = self._client
|
||||
if client is None or not client.is_ready():
|
||||
return None
|
||||
channel = self._known_channels.get(chat_id)
|
||||
if channel is not None:
|
||||
return channel
|
||||
channel_id = int(chat_id)
|
||||
channel = client.get_channel(channel_id)
|
||||
if channel is not None:
|
||||
return channel
|
||||
try:
|
||||
return await client.fetch_channel(channel_id)
|
||||
except Exception as e:
|
||||
self.logger.warning("channel {} unavailable: {}", chat_id, e)
|
||||
return None
|
||||
|
||||
async def _finalize_stream(self, chat_id: str, buf: _StreamBuf) -> None:
|
||||
"""Commit the final streamed content and flush overflow chunks."""
|
||||
chunks = DiscordBotClient._build_chunks(buf.text, [], False)
|
||||
if not chunks:
|
||||
self._stream_bufs.pop(chat_id, None)
|
||||
return
|
||||
|
||||
try:
|
||||
await buf.message.edit(content=chunks[0])
|
||||
except Exception as e:
|
||||
self.logger.warning("final stream edit failed: {}", e)
|
||||
raise
|
||||
|
||||
target = getattr(buf.message, "channel", None) or await self._resolve_channel(chat_id)
|
||||
if target is None:
|
||||
self.logger.warning("stream follow-up target {} unavailable", chat_id)
|
||||
self._stream_bufs.pop(chat_id, None)
|
||||
return
|
||||
|
||||
for extra_chunk in chunks[1:]:
|
||||
await target.send(content=extra_chunk)
|
||||
|
||||
self._stream_bufs.pop(chat_id, None)
|
||||
await self._stop_typing(chat_id)
|
||||
await self._clear_reactions(chat_id)
|
||||
|
||||
def _should_accept_inbound(
|
||||
self,
|
||||
message: discord.Message,
|
||||
sender_id: str,
|
||||
content: str,
|
||||
) -> bool:
|
||||
"""Check if inbound Discord message should be processed."""
|
||||
if not self.is_allowed(sender_id):
|
||||
return False
|
||||
# Channel-based filtering: only respond in allowed channels
|
||||
allow_channels = self.config.allow_channels
|
||||
if allow_channels:
|
||||
channel_ids = self._channel_allow_keys(message.channel)
|
||||
if channel_ids.isdisjoint(allow_channels):
|
||||
return False
|
||||
if message.guild is not None and not self._should_respond_in_group(message, content):
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _download_attachments(
|
||||
self,
|
||||
attachments: list[discord.Attachment],
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""Download supported attachments and return paths + display markers."""
|
||||
media_paths: list[str] = []
|
||||
markers: list[str] = []
|
||||
media_dir = get_media_dir("discord")
|
||||
|
||||
for attachment in attachments:
|
||||
filename = attachment.filename or "attachment"
|
||||
if attachment.size and attachment.size > MAX_ATTACHMENT_BYTES:
|
||||
markers.append(f"[attachment: {filename} - too large]")
|
||||
continue
|
||||
try:
|
||||
media_dir.mkdir(parents=True, exist_ok=True)
|
||||
safe_name = safe_filename(filename)
|
||||
file_path = media_dir / f"{attachment.id}_{safe_name}"
|
||||
await attachment.save(file_path)
|
||||
media_paths.append(str(file_path))
|
||||
markers.append(f"[attachment: {file_path.name}]")
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to download attachment: {}", e)
|
||||
markers.append(f"[attachment: {filename} - download failed]")
|
||||
|
||||
return media_paths, markers
|
||||
|
||||
@staticmethod
|
||||
def _compose_inbound_content(content: str, attachment_markers: list[str]) -> str:
|
||||
"""Combine message text with attachment markers."""
|
||||
content_parts = [content] if content else []
|
||||
content_parts.extend(attachment_markers)
|
||||
return "\n".join(part for part in content_parts if part) or "[empty message]"
|
||||
|
||||
@staticmethod
|
||||
def _is_system_message(message: discord.Message) -> bool:
|
||||
"""Return True for Discord system messages that carry no user prompt."""
|
||||
message_type = getattr(message, "type", discord.MessageType.default)
|
||||
return message_type not in {discord.MessageType.default, discord.MessageType.reply}
|
||||
|
||||
@staticmethod
|
||||
def _build_inbound_metadata(message: discord.Message) -> dict[str, str | None]:
|
||||
"""Build metadata for inbound Discord messages."""
|
||||
reply_to = (
|
||||
str(message.reference.message_id)
|
||||
if message.reference and message.reference.message_id
|
||||
else None
|
||||
)
|
||||
return {
|
||||
"message_id": str(message.id),
|
||||
"guild_id": str(message.guild.id) if message.guild else None,
|
||||
"reply_to": reply_to,
|
||||
}
|
||||
|
||||
def _should_respond_in_group(self, message: discord.Message, content: str) -> bool:
|
||||
"""Check if the bot should respond in a guild channel based on policy."""
|
||||
if self.config.group_policy == "open":
|
||||
return True
|
||||
|
||||
if self.config.group_policy == "mention":
|
||||
bot_user_id = self._bot_user_id
|
||||
if bot_user_id is None and self._client and self._client.user:
|
||||
bot_user_id = str(self._client.user.id)
|
||||
if bot_user_id is None:
|
||||
self.logger.debug(
|
||||
"message in {} ignored (bot identity unavailable)", message.channel.id
|
||||
)
|
||||
return False
|
||||
|
||||
if any(str(user.id) == bot_user_id for user in message.mentions):
|
||||
return True
|
||||
if bot_user_id in {str(user_id) for user_id in getattr(message, "raw_mentions", [])}:
|
||||
return True
|
||||
if f"<@{bot_user_id}>" in content or f"<@!{bot_user_id}>" in content:
|
||||
return True
|
||||
if self._references_bot_message(message, bot_user_id):
|
||||
return True
|
||||
|
||||
self.logger.debug("message in {} ignored (bot not mentioned)", message.channel.id)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _references_bot_message(message: discord.Message, bot_user_id: str) -> bool:
|
||||
"""Return True when a Discord reply targets a message authored by this bot."""
|
||||
reference = getattr(message, "reference", None)
|
||||
if reference is None:
|
||||
return False
|
||||
referenced_message = getattr(reference, "resolved", None) or getattr(
|
||||
reference, "cached_message", None
|
||||
)
|
||||
author = getattr(referenced_message, "author", None)
|
||||
return str(getattr(author, "id", "")) == bot_user_id
|
||||
|
||||
async def _start_typing(self, channel: Messageable) -> None:
|
||||
"""Start periodic typing indicator for a channel."""
|
||||
channel_id = self._channel_key(channel)
|
||||
await self._stop_typing(channel_id)
|
||||
|
||||
async def typing_loop() -> None:
|
||||
while self._running:
|
||||
try:
|
||||
async with channel.typing():
|
||||
await asyncio.sleep(TYPING_INTERVAL_S)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
except Exception as e:
|
||||
self.logger.debug("typing indicator failed for {}: {}", channel_id, e)
|
||||
return
|
||||
|
||||
self._typing_tasks[channel_id] = asyncio.create_task(typing_loop())
|
||||
|
||||
async def _stop_typing(self, channel_id: str) -> None:
|
||||
"""Stop typing indicator for a channel."""
|
||||
task = self._typing_tasks.pop(self._channel_key(channel_id), None)
|
||||
if task is None:
|
||||
return
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
async def _clear_reactions(self, chat_id: str) -> None:
|
||||
"""Remove all pending reactions after bot replies."""
|
||||
# Cancel delayed working emoji if it hasn't fired yet
|
||||
task = self._working_emoji_tasks.pop(chat_id, None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
|
||||
msg_obj = self._pending_reactions.pop(chat_id, None)
|
||||
if msg_obj is None:
|
||||
return
|
||||
bot_user = self._client.user if self._client else None
|
||||
for emoji in (self.config.read_receipt_emoji, self.config.working_emoji):
|
||||
with suppress(Exception):
|
||||
await msg_obj.remove_reaction(emoji, bot_user)
|
||||
|
||||
async def _cancel_all_typing(self) -> None:
|
||||
"""Stop all typing tasks."""
|
||||
channel_ids = list(self._typing_tasks)
|
||||
for channel_id in channel_ids:
|
||||
await self._stop_typing(channel_id)
|
||||
|
||||
async def _reset_runtime_state(self, close_client: bool) -> None:
|
||||
"""Reset client and typing state."""
|
||||
await self._cancel_all_typing()
|
||||
self._stream_bufs.clear()
|
||||
self._known_channels.clear()
|
||||
if close_client and self._client is not None and not self._client.is_closed():
|
||||
try:
|
||||
await self._client.close()
|
||||
except Exception as e:
|
||||
self.logger.warning("client close failed: {}", e)
|
||||
self._client = None
|
||||
self._bot_user_id = None
|
||||
@@ -0,0 +1,914 @@
|
||||
"""Email channel implementation using IMAP polling + SMTP replies."""
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
import imaplib
|
||||
import mimetypes
|
||||
import re
|
||||
import smtplib
|
||||
import ssl
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from datetime import date
|
||||
from email import policy
|
||||
from email.header import decode_header, make_header
|
||||
from email.message import EmailMessage
|
||||
from email.parser import BytesParser
|
||||
from email.utils import parseaddr
|
||||
from fnmatch import fnmatch
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import logging; logger = logging.getLogger(__name__)
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_media_dir
|
||||
from pydantic import BaseModel
|
||||
from src.channels.utils import safe_filename
|
||||
|
||||
|
||||
class EmailConfig(BaseModel):
|
||||
"""Email channel configuration (IMAP inbound + SMTP outbound)."""
|
||||
|
||||
enabled: bool = False
|
||||
consent_granted: bool = False
|
||||
|
||||
imap_host: str = ""
|
||||
imap_port: int = 993
|
||||
imap_username: str = ""
|
||||
imap_password: str = ""
|
||||
imap_mailbox: str = "INBOX"
|
||||
imap_use_ssl: bool = True
|
||||
|
||||
smtp_host: str = ""
|
||||
smtp_port: int = 587
|
||||
smtp_username: str = ""
|
||||
smtp_password: str = ""
|
||||
smtp_use_tls: bool = True
|
||||
smtp_use_ssl: bool = False
|
||||
from_address: str = ""
|
||||
|
||||
auto_reply_enabled: bool = True
|
||||
poll_interval_seconds: int = 30
|
||||
mark_seen: bool = True
|
||||
post_action: Literal["delete", "move"] | None = None
|
||||
post_action_move_mailbox: str | None = None
|
||||
post_action_expunge: bool = False
|
||||
post_action_ignore_skipped: bool = True
|
||||
max_body_chars: int = 12000
|
||||
subject_prefix: str = "Re: "
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
|
||||
# Email authentication verification (anti-spoofing)
|
||||
verify_dkim: bool = True # Require Authentication-Results with dkim=pass
|
||||
verify_spf: bool = True # Require Authentication-Results with spf=pass
|
||||
|
||||
# Attachment handling — set allowed types to enable (e.g. ["application/pdf", "image/*"], or ["*"] for all)
|
||||
allowed_attachment_types: list[str] = Field(default_factory=list)
|
||||
max_attachment_size: int = 2_000_000 # 2MB per attachment
|
||||
max_attachments_per_email: int = 5
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ServerFeatures:
|
||||
move: bool
|
||||
uidplus: bool
|
||||
uid_store: bool | None = None
|
||||
|
||||
|
||||
class EmailChannel(BaseChannel):
|
||||
"""
|
||||
Email channel.
|
||||
|
||||
Inbound:
|
||||
- Poll IMAP mailbox for unread messages.
|
||||
- Convert each message into an inbound event.
|
||||
|
||||
Outbound:
|
||||
- Send responses via SMTP back to the sender address.
|
||||
"""
|
||||
|
||||
name = "email"
|
||||
display_name = "Email"
|
||||
_IMAP_MONTHS = (
|
||||
"Jan",
|
||||
"Feb",
|
||||
"Mar",
|
||||
"Apr",
|
||||
"May",
|
||||
"Jun",
|
||||
"Jul",
|
||||
"Aug",
|
||||
"Sep",
|
||||
"Oct",
|
||||
"Nov",
|
||||
"Dec",
|
||||
)
|
||||
_IMAP_RECONNECT_MARKERS = (
|
||||
"disconnected for inactivity",
|
||||
"eof occurred in violation of protocol",
|
||||
"socket error",
|
||||
"connection reset",
|
||||
"broken pipe",
|
||||
"bye",
|
||||
)
|
||||
_IMAP_MISSING_MAILBOX_MARKERS = (
|
||||
"mailbox doesn't exist",
|
||||
"select failed",
|
||||
"no such mailbox",
|
||||
"can't open mailbox",
|
||||
"does not exist",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return EmailConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = EmailConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: EmailConfig = config
|
||||
self._self_addresses = self._collect_self_addresses()
|
||||
self._last_subject_by_chat: dict[str, str] = {}
|
||||
self._last_message_id_by_chat: dict[str, str] = {}
|
||||
self._processed_uids: set[str] = set() # Capped to prevent unbounded growth
|
||||
self._MAX_PROCESSED_UIDS = 100000
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start polling IMAP for inbound emails."""
|
||||
if not self.config.consent_granted:
|
||||
self.logger.warning(
|
||||
"Email channel disabled: consent_granted is false. "
|
||||
"Set channels.email.consentGranted=true after explicit user permission."
|
||||
)
|
||||
return
|
||||
|
||||
if not self._validate_config():
|
||||
return
|
||||
|
||||
self._running = True
|
||||
if not self.config.verify_dkim and not self.config.verify_spf:
|
||||
self.logger.warning(
|
||||
"DKIM and SPF verification are both DISABLED. "
|
||||
"Emails with spoofed From headers will be accepted. "
|
||||
"Set verify_dkim=true and verify_spf=true for anti-spoofing protection."
|
||||
)
|
||||
self.logger.info("Starting Email channel (IMAP polling mode)...")
|
||||
|
||||
poll_seconds = max(5, int(self.config.poll_interval_seconds))
|
||||
while self._running:
|
||||
try:
|
||||
inbound_items, skipped_uids = await asyncio.to_thread(self._fetch_new_messages)
|
||||
should_apply_post_action = self._should_apply_post_action()
|
||||
post_actions_uids: set[str] = set()
|
||||
for item in inbound_items:
|
||||
sender = item["sender"]
|
||||
subject = item.get("subject", "")
|
||||
message_id = item.get("message_id", "")
|
||||
|
||||
if subject:
|
||||
self._last_subject_by_chat[sender] = subject
|
||||
if message_id:
|
||||
self._last_message_id_by_chat[sender] = message_id
|
||||
|
||||
try:
|
||||
await self._handle_message(
|
||||
sender_id=sender,
|
||||
chat_id=sender,
|
||||
content=item["content"],
|
||||
media=item.get("media") or None,
|
||||
metadata=item.get("metadata", {}),
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Error delivering email from {}", sender)
|
||||
continue
|
||||
|
||||
uid = str((item.get("metadata") or {}).get("uid") or "")
|
||||
if uid and should_apply_post_action:
|
||||
post_actions_uids.add(uid)
|
||||
|
||||
if should_apply_post_action and not self.config.post_action_ignore_skipped:
|
||||
post_actions_uids.update(skipped_uids)
|
||||
|
||||
if post_actions_uids:
|
||||
await asyncio.to_thread(self._apply_post_actions_batch, sorted(post_actions_uids))
|
||||
except Exception:
|
||||
self.logger.exception("Polling error")
|
||||
|
||||
if not self._running:
|
||||
break
|
||||
await asyncio.sleep(poll_seconds)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop polling loop."""
|
||||
self._running = False
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send email via SMTP."""
|
||||
if not self.config.consent_granted:
|
||||
self.logger.warning("Skip email send: consent_granted is false")
|
||||
return
|
||||
|
||||
if not self.config.smtp_host:
|
||||
self.logger.warning("SMTP host not configured")
|
||||
return
|
||||
|
||||
# Skip progress messages to prevent sending an empty email after each tool call
|
||||
if (msg.metadata or {}).get("_progress"):
|
||||
self.logger.debug("Skip progress message to {}", msg.chat_id)
|
||||
return
|
||||
|
||||
to_addr = msg.chat_id.strip()
|
||||
if not to_addr:
|
||||
self.logger.warning("Missing recipient address")
|
||||
return
|
||||
|
||||
# Determine if this is a reply (recipient has sent us an email before)
|
||||
is_reply = to_addr in self._last_subject_by_chat
|
||||
force_send = bool((msg.metadata or {}).get("force_send"))
|
||||
|
||||
# autoReplyEnabled only controls automatic replies, not proactive sends
|
||||
if is_reply and not self.config.auto_reply_enabled and not force_send:
|
||||
self.logger.info("Skip automatic reply to {}: auto_reply_enabled is false", to_addr)
|
||||
return
|
||||
|
||||
base_subject = self._last_subject_by_chat.get(to_addr, "vibe-trading reply")
|
||||
subject = self._reply_subject(base_subject)
|
||||
if msg.metadata and isinstance(msg.metadata.get("subject"), str):
|
||||
override = msg.metadata["subject"].strip()
|
||||
if override:
|
||||
subject = override
|
||||
|
||||
attachments: list[tuple[bytes, str, str, str]] = []
|
||||
failed_attachments: list[str] = []
|
||||
max_attachment_size = max(0, int(self.config.max_attachment_size))
|
||||
max_attachment_count = max(0, int(self.config.max_attachments_per_email))
|
||||
for media_path in msg.media or []:
|
||||
path = Path(media_path)
|
||||
filename = path.name or "attachment"
|
||||
if len(attachments) >= max_attachment_count:
|
||||
failed_attachments.append(f"[attachment: {filename} - too many attachments]")
|
||||
self.logger.warning("Attachment count limit reached, skipping: {}", media_path)
|
||||
continue
|
||||
if not path.is_file():
|
||||
failed_attachments.append(f"[attachment: {filename} - send failed]")
|
||||
self.logger.warning("Attachment not found, skipping: {}", media_path)
|
||||
continue
|
||||
try:
|
||||
size = path.stat().st_size
|
||||
if max_attachment_size <= 0 or size > max_attachment_size:
|
||||
failed_attachments.append(f"[attachment: {filename} - too large]")
|
||||
self.logger.warning(
|
||||
"Attachment too large, skipping: {} ({} > {} bytes)",
|
||||
media_path,
|
||||
size,
|
||||
max_attachment_size,
|
||||
)
|
||||
continue
|
||||
data = path.read_bytes()
|
||||
ctype, _ = mimetypes.guess_type(str(path))
|
||||
if ctype is None:
|
||||
ctype = "application/octet-stream"
|
||||
maintype, subtype = ctype.split("/", 1)
|
||||
attachments.append((data, maintype, subtype, filename))
|
||||
self.logger.info("Attached file: {}", filename)
|
||||
except Exception:
|
||||
failed_attachments.append(f"[attachment: {filename} - send failed]")
|
||||
self.logger.exception("Failed to attach file {}", media_path)
|
||||
|
||||
content = msg.content or ""
|
||||
if failed_attachments:
|
||||
fallback = "\n".join(failed_attachments)
|
||||
content = f"{content.rstrip()}\n\n{fallback}" if content.strip() else fallback
|
||||
|
||||
email_msg = EmailMessage()
|
||||
email_msg["From"] = self.config.from_address or self.config.smtp_username or self.config.imap_username
|
||||
email_msg["To"] = to_addr
|
||||
email_msg["Subject"] = subject
|
||||
email_msg.set_content(content)
|
||||
|
||||
for data, maintype, subtype, filename in attachments:
|
||||
email_msg.add_attachment(
|
||||
data,
|
||||
maintype=maintype,
|
||||
subtype=subtype,
|
||||
filename=filename,
|
||||
)
|
||||
|
||||
in_reply_to = self._last_message_id_by_chat.get(to_addr)
|
||||
if in_reply_to:
|
||||
email_msg["In-Reply-To"] = in_reply_to
|
||||
email_msg["References"] = in_reply_to
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(self._smtp_send, email_msg)
|
||||
except Exception:
|
||||
self.logger.exception("Error sending to {}", to_addr)
|
||||
raise
|
||||
|
||||
def _validate_config(self) -> bool:
|
||||
missing = []
|
||||
if not self.config.imap_host:
|
||||
missing.append("imap_host")
|
||||
if not self.config.imap_username:
|
||||
missing.append("imap_username")
|
||||
if not self.config.imap_password:
|
||||
missing.append("imap_password")
|
||||
if not self.config.smtp_host:
|
||||
missing.append("smtp_host")
|
||||
if not self.config.smtp_username:
|
||||
missing.append("smtp_username")
|
||||
if not self.config.smtp_password:
|
||||
missing.append("smtp_password")
|
||||
|
||||
if self.config.post_action == "move" and not (self.config.post_action_move_mailbox or "").strip():
|
||||
missing.append("post_action_move_mailbox")
|
||||
|
||||
if missing:
|
||||
self.logger.error("Channel not configured, missing: {}", ', '.join(missing))
|
||||
return False
|
||||
return True
|
||||
|
||||
def _smtp_send(self, msg: EmailMessage) -> None:
|
||||
timeout = 30
|
||||
if self.config.smtp_use_ssl:
|
||||
with smtplib.SMTP_SSL(
|
||||
self.config.smtp_host,
|
||||
self.config.smtp_port,
|
||||
timeout=timeout,
|
||||
) as smtp:
|
||||
smtp.login(self.config.smtp_username, self.config.smtp_password)
|
||||
smtp.send_message(msg)
|
||||
return
|
||||
|
||||
with smtplib.SMTP(self.config.smtp_host, self.config.smtp_port, timeout=timeout) as smtp:
|
||||
if self.config.smtp_use_tls:
|
||||
smtp.starttls(context=ssl.create_default_context())
|
||||
smtp.login(self.config.smtp_username, self.config.smtp_password)
|
||||
smtp.send_message(msg)
|
||||
|
||||
def _fetch_new_messages(self) -> tuple[list[dict[str, Any]], set[str]]:
|
||||
"""Poll IMAP and return parsed unread messages plus skipped message UIDs."""
|
||||
return self._fetch_messages(
|
||||
search_criteria=("UNSEEN",),
|
||||
mark_seen=self.config.mark_seen,
|
||||
dedupe=True,
|
||||
limit=0,
|
||||
)
|
||||
|
||||
def fetch_messages_between_dates(
|
||||
self,
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
limit: int = 20,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Fetch messages in [start_date, end_date) by IMAP date search.
|
||||
|
||||
This is used for historical summarization tasks (e.g. "yesterday").
|
||||
"""
|
||||
if end_date <= start_date:
|
||||
return []
|
||||
|
||||
messages, _ = self._fetch_messages(
|
||||
search_criteria=(
|
||||
"SINCE",
|
||||
self._format_imap_date(start_date),
|
||||
"BEFORE",
|
||||
self._format_imap_date(end_date),
|
||||
),
|
||||
mark_seen=False,
|
||||
dedupe=False,
|
||||
limit=max(1, int(limit)),
|
||||
)
|
||||
return messages
|
||||
|
||||
def _fetch_messages(
|
||||
self,
|
||||
search_criteria: tuple[str, ...],
|
||||
mark_seen: bool,
|
||||
dedupe: bool,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], set[str]]:
|
||||
messages: list[dict[str, Any]] = []
|
||||
skipped_uids: set[str] = set()
|
||||
cycle_uids: set[str] = set()
|
||||
|
||||
for attempt in range(2):
|
||||
try:
|
||||
self._fetch_messages_once(
|
||||
search_criteria,
|
||||
mark_seen,
|
||||
dedupe,
|
||||
limit,
|
||||
messages,
|
||||
skipped_uids,
|
||||
cycle_uids,
|
||||
)
|
||||
return messages, skipped_uids
|
||||
except Exception as exc:
|
||||
if attempt == 1 or not self._is_stale_imap_error(exc):
|
||||
raise
|
||||
self.logger.warning("IMAP connection went stale, retrying once: {}", exc)
|
||||
|
||||
return messages, skipped_uids
|
||||
|
||||
def _fetch_messages_once(
|
||||
self,
|
||||
search_criteria: tuple[str, ...],
|
||||
mark_seen: bool,
|
||||
dedupe: bool,
|
||||
limit: int,
|
||||
messages: list[dict[str, Any]],
|
||||
skipped_uids: set[str],
|
||||
cycle_uids: set[str],
|
||||
) -> None:
|
||||
"""Fetch messages by arbitrary IMAP search criteria."""
|
||||
mailbox = self.config.imap_mailbox or "INBOX"
|
||||
|
||||
client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True)
|
||||
if client is None:
|
||||
return messages
|
||||
|
||||
try:
|
||||
status, data = client.search(None, *search_criteria)
|
||||
if status != "OK" or not data:
|
||||
return messages
|
||||
|
||||
ids = data[0].split()
|
||||
if limit > 0 and len(ids) > limit:
|
||||
ids = ids[-limit:]
|
||||
for imap_id in ids:
|
||||
status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)")
|
||||
if status != "OK" or not fetched:
|
||||
continue
|
||||
|
||||
raw_bytes = self._extract_message_bytes(fetched)
|
||||
if raw_bytes is None:
|
||||
continue
|
||||
|
||||
uid = self._extract_uid(fetched)
|
||||
if uid and uid in cycle_uids:
|
||||
continue
|
||||
if dedupe and uid and uid in self._processed_uids:
|
||||
continue
|
||||
|
||||
parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
|
||||
sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
|
||||
if not sender:
|
||||
continue
|
||||
if self._is_self_address(sender):
|
||||
self.logger.info("From {} ignored: matches bot-owned address", sender)
|
||||
self._remember_processed_uid(uid, dedupe, cycle_uids)
|
||||
if mark_seen:
|
||||
client.store(imap_id, "+FLAGS", "\\Seen")
|
||||
if uid:
|
||||
skipped_uids.add(uid)
|
||||
continue
|
||||
|
||||
# --- Anti-spoofing: verify Authentication-Results ---
|
||||
spf_pass, dkim_pass = self._check_authentication_results(parsed)
|
||||
if self.config.verify_spf and not spf_pass:
|
||||
self.logger.warning(
|
||||
"From {} rejected: SPF verification failed "
|
||||
"(no 'spf=pass' in Authentication-Results header)",
|
||||
sender,
|
||||
)
|
||||
self._remember_processed_uid(uid, dedupe, cycle_uids)
|
||||
if uid:
|
||||
skipped_uids.add(uid)
|
||||
continue
|
||||
if self.config.verify_dkim and not dkim_pass:
|
||||
self.logger.warning(
|
||||
"From {} rejected: DKIM verification failed "
|
||||
"(no 'dkim=pass' in Authentication-Results header)",
|
||||
sender,
|
||||
)
|
||||
self._remember_processed_uid(uid, dedupe, cycle_uids)
|
||||
if uid:
|
||||
skipped_uids.add(uid)
|
||||
continue
|
||||
|
||||
if not self.is_allowed(sender):
|
||||
self._remember_processed_uid(uid, dedupe, cycle_uids)
|
||||
if mark_seen:
|
||||
client.store(imap_id, "+FLAGS", "\\Seen")
|
||||
if uid:
|
||||
skipped_uids.add(uid)
|
||||
continue
|
||||
|
||||
subject = self._decode_header_value(parsed.get("Subject", ""))
|
||||
date_value = parsed.get("Date", "")
|
||||
message_id = parsed.get("Message-ID", "").strip()
|
||||
body = self._extract_text_body(parsed)
|
||||
|
||||
if not body:
|
||||
body = "(empty email body)"
|
||||
|
||||
body = body[: self.config.max_body_chars]
|
||||
content = (
|
||||
f"[EMAIL-CONTEXT] Email received.\n"
|
||||
f"From: {sender}\n"
|
||||
f"Subject: {subject}\n"
|
||||
f"Date: {date_value}\n\n"
|
||||
f"{body}"
|
||||
)
|
||||
|
||||
# --- Attachment extraction ---
|
||||
attachment_paths: list[str] = []
|
||||
if self.config.allowed_attachment_types:
|
||||
saved = self._extract_attachments(
|
||||
parsed,
|
||||
uid or "noid",
|
||||
allowed_types=self.config.allowed_attachment_types,
|
||||
max_size=self.config.max_attachment_size,
|
||||
max_count=self.config.max_attachments_per_email,
|
||||
)
|
||||
for p in saved:
|
||||
attachment_paths.append(str(p))
|
||||
content += f"\n[attachment: {p.name} — saved to {p}]"
|
||||
|
||||
metadata = {
|
||||
"message_id": message_id,
|
||||
"subject": subject,
|
||||
"date": date_value,
|
||||
"sender_email": sender,
|
||||
"uid": uid,
|
||||
}
|
||||
messages.append(
|
||||
{
|
||||
"sender": sender,
|
||||
"subject": subject,
|
||||
"message_id": message_id,
|
||||
"content": content,
|
||||
"metadata": metadata,
|
||||
"media": attachment_paths,
|
||||
}
|
||||
)
|
||||
|
||||
self._remember_processed_uid(uid, dedupe, cycle_uids)
|
||||
|
||||
if mark_seen:
|
||||
client.store(imap_id, "+FLAGS", "\\Seen")
|
||||
finally:
|
||||
self._close_imap_client(client)
|
||||
|
||||
def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None:
|
||||
if self.config.imap_use_ssl:
|
||||
client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
|
||||
else:
|
||||
client = imaplib.IMAP4(self.config.imap_host, self.config.imap_port)
|
||||
|
||||
try:
|
||||
client.login(self.config.imap_username, self.config.imap_password)
|
||||
try:
|
||||
status, _ = client.select(mailbox)
|
||||
except Exception as exc:
|
||||
if missing_mailbox_ok and self._is_missing_mailbox_error(exc):
|
||||
self.logger.warning("Mailbox unavailable, skipping poll for {}: {}", mailbox, exc)
|
||||
self._close_imap_client(client)
|
||||
return None
|
||||
raise
|
||||
|
||||
if status != "OK":
|
||||
self.logger.warning("Mailbox select returned {}, skipping poll for {}", status, mailbox)
|
||||
self._close_imap_client(client)
|
||||
return None
|
||||
except Exception:
|
||||
self._close_imap_client(client)
|
||||
raise
|
||||
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def _close_imap_client(client: Any) -> None:
|
||||
with suppress(Exception):
|
||||
client.logout()
|
||||
|
||||
def _collect_self_addresses(self) -> set[str]:
|
||||
"""Return normalized email addresses owned by this channel instance."""
|
||||
candidates = (
|
||||
self.config.from_address,
|
||||
self.config.smtp_username,
|
||||
self.config.imap_username,
|
||||
)
|
||||
normalized = {
|
||||
addr
|
||||
for candidate in candidates
|
||||
if (addr := self._normalize_address(candidate))
|
||||
}
|
||||
return normalized
|
||||
|
||||
@staticmethod
|
||||
def _normalize_address(value: str) -> str:
|
||||
"""Normalize an address or mailbox-like identifier for comparisons."""
|
||||
raw = (value or "").strip()
|
||||
if not raw:
|
||||
return ""
|
||||
parsed = parseaddr(raw)[1].strip().lower()
|
||||
if parsed:
|
||||
return parsed
|
||||
if "@" in raw:
|
||||
return raw.lower()
|
||||
return ""
|
||||
|
||||
def _is_self_address(self, sender: str) -> bool:
|
||||
"""Return True when an inbound sender belongs to the bot itself."""
|
||||
normalized_sender = self._normalize_address(sender)
|
||||
return bool(normalized_sender) and normalized_sender in self._self_addresses
|
||||
|
||||
def _remember_processed_uid(self, uid: str, dedupe: bool, cycle_uids: set[str]) -> None:
|
||||
"""Track a fetched UID so skipped messages are not reprocessed forever."""
|
||||
if not uid:
|
||||
return
|
||||
cycle_uids.add(uid)
|
||||
if dedupe:
|
||||
self._processed_uids.add(uid)
|
||||
# mark_seen is the primary dedup; this set is a safety net
|
||||
if len(self._processed_uids) > self._MAX_PROCESSED_UIDS:
|
||||
# Evict a random half to cap memory; mark_seen is the primary dedup
|
||||
self._processed_uids = set(list(self._processed_uids)[len(self._processed_uids) // 2:])
|
||||
|
||||
def _should_apply_post_action(self) -> bool:
|
||||
return self.config.post_action in {"delete", "move"}
|
||||
|
||||
def _apply_post_actions_batch(self, post_actions_uids: list[str]) -> None:
|
||||
if not self._should_apply_post_action() or not post_actions_uids:
|
||||
return
|
||||
|
||||
mailbox = self.config.imap_mailbox or "INBOX"
|
||||
client = self._open_imap_client(mailbox=mailbox)
|
||||
if client is None:
|
||||
return
|
||||
|
||||
try:
|
||||
features = self._server_features(client)
|
||||
# Apply all post-actions in one IMAP session. `features` also carries
|
||||
# session-learned behavior (e.g. UID STORE support) so later UIDs can
|
||||
# skip known-broken paths.
|
||||
for uid in post_actions_uids:
|
||||
if uid:
|
||||
self._apply_post_action(client, uid, features)
|
||||
finally:
|
||||
self._close_imap_client(client)
|
||||
|
||||
def _apply_post_action(
|
||||
self,
|
||||
client: Any,
|
||||
uid: str,
|
||||
features: _ServerFeatures,
|
||||
) -> None:
|
||||
action = self.config.post_action
|
||||
|
||||
if action == "delete":
|
||||
if not self._uid_store_deleted(client, uid, features):
|
||||
return
|
||||
self._uid_expunge_or_fallback(client, uid, features)
|
||||
return
|
||||
|
||||
if action == "move":
|
||||
target = (self.config.post_action_move_mailbox or "").strip()
|
||||
if features.move:
|
||||
status, _ = client.uid("MOVE", uid, target)
|
||||
if status != "OK":
|
||||
self.logger.warning("Post-action move failed (UID MOVE) for UID {} to mailbox {}", uid, target)
|
||||
return
|
||||
|
||||
status, _ = client.uid("COPY", uid, target)
|
||||
if status != "OK":
|
||||
self.logger.warning("Post-action move failed (UID COPY) for UID {} to mailbox {}", uid, target)
|
||||
return
|
||||
if not self._uid_store_deleted(client, uid, features):
|
||||
return
|
||||
self._uid_expunge_or_fallback(client, uid, features)
|
||||
|
||||
@staticmethod
|
||||
def _server_features(client: Any) -> _ServerFeatures:
|
||||
caps: set[str] = set()
|
||||
with suppress(Exception):
|
||||
status, data = client.capability()
|
||||
if status == "OK" and data:
|
||||
for raw in data:
|
||||
if isinstance(raw, (bytes, bytearray)):
|
||||
caps.update(token.upper() for token in raw.decode("utf-8", errors="ignore").split())
|
||||
elif isinstance(raw, str):
|
||||
caps.update(token.upper() for token in raw.split())
|
||||
return _ServerFeatures(move="MOVE" in caps, uidplus="UIDPLUS" in caps)
|
||||
|
||||
@staticmethod
|
||||
def _lookup_imap_id_by_uid(client: Any, uid: str) -> bytes | None:
|
||||
# IMAP exposes two message identifiers: UID (stable) and sequence number
|
||||
# (session-local). We target by UID first, but some servers may reject
|
||||
# UID STORE. In that case we resolve the current sequence number for the
|
||||
# UID and retry with STORE using that sequence id.
|
||||
status, data = client.search(None, "UID", uid)
|
||||
if status != "OK" or not data or not data[0]:
|
||||
return None
|
||||
return data[0].split()[0]
|
||||
|
||||
def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool:
|
||||
# Optimistic path: try UID STORE first because UID is stable and avoids
|
||||
# sequence-number lookup. If this fails once for the session, remember it
|
||||
# and use the sequence STORE fallback directly for remaining UIDs.
|
||||
if features.uid_store is not False:
|
||||
status, _ = client.uid("STORE", uid, "+FLAGS", "(\\Deleted)")
|
||||
if status == "OK":
|
||||
features.uid_store = True
|
||||
return True
|
||||
features.uid_store = False
|
||||
|
||||
# Compatibility fallback for servers where UID STORE is unavailable or
|
||||
# unreliable: resolve the current sequence number from UID and use STORE.
|
||||
imap_id = self._lookup_imap_id_by_uid(client, uid)
|
||||
if not imap_id:
|
||||
self.logger.warning("Post-action skipped: UID {} not found", uid)
|
||||
return False
|
||||
|
||||
status, _ = client.store(imap_id, "+FLAGS", "\\Deleted")
|
||||
if status != "OK":
|
||||
self.logger.warning("Post-action failed: could not mark UID {} as deleted", uid)
|
||||
return False
|
||||
return True
|
||||
|
||||
def _uid_expunge_or_fallback(self, client: Any, uid: str, features: _ServerFeatures) -> None:
|
||||
# Prefer UID-scoped expunge when supported to avoid expunging unrelated
|
||||
# messages already marked \Deleted in the selected mailbox.
|
||||
if features.uidplus:
|
||||
status, _ = client.uid("EXPUNGE", uid)
|
||||
if status == "OK":
|
||||
return
|
||||
self.logger.warning("UID EXPUNGE failed for UID {}, falling back to EXPUNGE", uid)
|
||||
if self.config.post_action_expunge:
|
||||
client.expunge()
|
||||
|
||||
@classmethod
|
||||
def _is_stale_imap_error(cls, exc: Exception) -> bool:
|
||||
message = str(exc).lower()
|
||||
return any(marker in message for marker in cls._IMAP_RECONNECT_MARKERS)
|
||||
|
||||
@classmethod
|
||||
def _is_missing_mailbox_error(cls, exc: Exception) -> bool:
|
||||
message = str(exc).lower()
|
||||
return any(marker in message for marker in cls._IMAP_MISSING_MAILBOX_MARKERS)
|
||||
|
||||
@classmethod
|
||||
def _format_imap_date(cls, value: date) -> str:
|
||||
"""Format date for IMAP search (always English month abbreviations)."""
|
||||
month = cls._IMAP_MONTHS[value.month - 1]
|
||||
return f"{value.day:02d}-{month}-{value.year}"
|
||||
|
||||
@staticmethod
|
||||
def _extract_message_bytes(fetched: list[Any]) -> bytes | None:
|
||||
for item in fetched:
|
||||
if isinstance(item, tuple) and len(item) >= 2 and isinstance(item[1], (bytes, bytearray)):
|
||||
return bytes(item[1])
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _extract_uid(fetched: list[Any]) -> str:
|
||||
for item in fetched:
|
||||
if isinstance(item, tuple) and item and isinstance(item[0], (bytes, bytearray)):
|
||||
head = bytes(item[0]).decode("utf-8", errors="ignore")
|
||||
m = re.search(r"UID\s+(\d+)", head)
|
||||
if m:
|
||||
return m.group(1)
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _decode_header_value(value: str) -> str:
|
||||
if not value:
|
||||
return ""
|
||||
try:
|
||||
return str(make_header(decode_header(value)))
|
||||
except Exception:
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def _extract_text_body(cls, msg: Any) -> str:
|
||||
"""Best-effort extraction of readable body text."""
|
||||
if msg.is_multipart():
|
||||
plain_parts: list[str] = []
|
||||
html_parts: list[str] = []
|
||||
for part in msg.walk():
|
||||
if part.get_content_disposition() == "attachment":
|
||||
continue
|
||||
content_type = part.get_content_type()
|
||||
try:
|
||||
payload = part.get_content()
|
||||
except Exception:
|
||||
payload_bytes = part.get_payload(decode=True) or b""
|
||||
charset = part.get_content_charset() or "utf-8"
|
||||
payload = payload_bytes.decode(charset, errors="replace")
|
||||
if not isinstance(payload, str):
|
||||
continue
|
||||
if content_type == "text/plain":
|
||||
plain_parts.append(payload)
|
||||
elif content_type == "text/html":
|
||||
html_parts.append(payload)
|
||||
if plain_parts:
|
||||
return "\n\n".join(plain_parts).strip()
|
||||
if html_parts:
|
||||
return cls._html_to_text("\n\n".join(html_parts)).strip()
|
||||
return ""
|
||||
|
||||
try:
|
||||
payload = msg.get_content()
|
||||
except Exception:
|
||||
payload_bytes = msg.get_payload(decode=True) or b""
|
||||
charset = msg.get_content_charset() or "utf-8"
|
||||
payload = payload_bytes.decode(charset, errors="replace")
|
||||
if not isinstance(payload, str):
|
||||
return ""
|
||||
if msg.get_content_type() == "text/html":
|
||||
return cls._html_to_text(payload).strip()
|
||||
return payload.strip()
|
||||
|
||||
@staticmethod
|
||||
def _check_authentication_results(parsed_msg: Any) -> tuple[bool, bool]:
|
||||
"""Parse Authentication-Results headers for SPF and DKIM verdicts.
|
||||
|
||||
Returns:
|
||||
A tuple of (spf_pass, dkim_pass) booleans.
|
||||
"""
|
||||
spf_pass = False
|
||||
dkim_pass = False
|
||||
for ar_header in parsed_msg.get_all("Authentication-Results") or []:
|
||||
ar_lower = ar_header.lower()
|
||||
if re.search(r"\bspf\s*=\s*pass\b", ar_lower):
|
||||
spf_pass = True
|
||||
if re.search(r"\bdkim\s*=\s*pass\b", ar_lower):
|
||||
dkim_pass = True
|
||||
return spf_pass, dkim_pass
|
||||
|
||||
@classmethod
|
||||
def _extract_attachments(
|
||||
cls,
|
||||
msg: Any,
|
||||
uid: str,
|
||||
*,
|
||||
allowed_types: list[str],
|
||||
max_size: int,
|
||||
max_count: int,
|
||||
) -> list[Path]:
|
||||
"""Extract and save email attachments to the media directory.
|
||||
|
||||
Returns list of saved file paths.
|
||||
"""
|
||||
if not msg.is_multipart():
|
||||
return []
|
||||
|
||||
saved: list[Path] = []
|
||||
media_dir = get_media_dir("email")
|
||||
|
||||
for part in msg.walk():
|
||||
if len(saved) >= max_count:
|
||||
break
|
||||
if part.get_content_disposition() != "attachment":
|
||||
continue
|
||||
|
||||
content_type = part.get_content_type()
|
||||
if not any(fnmatch(content_type, pat) for pat in allowed_types):
|
||||
logger.debug("Attachment skipped (type {}): not in allowed list", content_type)
|
||||
continue
|
||||
|
||||
payload = part.get_payload(decode=True)
|
||||
if payload is None:
|
||||
continue
|
||||
if len(payload) > max_size:
|
||||
logger.warning(
|
||||
"Attachment skipped: size {} exceeds limit {}",
|
||||
len(payload),
|
||||
max_size,
|
||||
)
|
||||
continue
|
||||
|
||||
raw_name = part.get_filename() or "attachment"
|
||||
sanitized = safe_filename(raw_name) or "attachment"
|
||||
dest = media_dir / f"{uid}_{sanitized}"
|
||||
|
||||
try:
|
||||
dest.write_bytes(payload)
|
||||
saved.append(dest)
|
||||
logger.info("Attachment saved: {}", dest)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to save attachment {}: {}", dest, exc)
|
||||
|
||||
return saved
|
||||
|
||||
@staticmethod
|
||||
def _html_to_text(raw_html: str) -> str:
|
||||
text = re.sub(r"<\s*br\s*/?>", "\n", raw_html, flags=re.IGNORECASE)
|
||||
text = re.sub(r"<\s*/\s*p\s*>", "\n", text, flags=re.IGNORECASE)
|
||||
text = re.sub(r"<[^>]+>", "", text)
|
||||
return html.unescape(text)
|
||||
|
||||
def _reply_subject(self, base_subject: str) -> str:
|
||||
subject = (base_subject or "").strip() or "vibe-trading reply"
|
||||
prefix = self.config.subject_prefix or "Re: "
|
||||
if subject.lower().startswith("re:"):
|
||||
return subject
|
||||
return f"{prefix}{subject}"
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,478 @@
|
||||
"""Channel manager for coordinating chat channels."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.registry import (
|
||||
discover_channel_names,
|
||||
discover_plugins,
|
||||
inspect_channels,
|
||||
load_channel_class,
|
||||
)
|
||||
from src.config.paths import get_workspace_path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Retry delays for message sending (exponential backoff: 1s, 2s, 4s)
|
||||
_SEND_RETRY_DELAYS = (1, 2, 4)
|
||||
|
||||
_BOOL_CAMEL_ALIASES: dict[str, str] = {
|
||||
"send_progress": "sendProgress",
|
||||
"send_tool_hints": "sendToolHints",
|
||||
"show_reasoning": "showReasoning",
|
||||
}
|
||||
|
||||
|
||||
class ChannelManager:
|
||||
"""Manages chat channels and coordinates message routing.
|
||||
|
||||
Responsibilities:
|
||||
- Initialize enabled channels (Telegram, Discord, Slack, etc.)
|
||||
- Start/stop channels
|
||||
- Route outbound messages
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Any, # ChannelsConfig or dict
|
||||
bus: MessageBus,
|
||||
*,
|
||||
session_service: Any | None = None,
|
||||
cron_service: Any | None = None,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.bus = bus
|
||||
self._session_service = session_service
|
||||
self._cron_service = cron_service
|
||||
self.channels: dict[str, BaseChannel] = {}
|
||||
self._dispatch_task: asyncio.Task | None = None
|
||||
self._origin_reply_fingerprints: dict[tuple[str, str, str], str] = {}
|
||||
self._status: dict[str, dict[str, Any]] = {}
|
||||
|
||||
self._init_channels()
|
||||
|
||||
def _init_channels(self) -> None:
|
||||
"""Initialize channels discovered via pkgutil scan + entry_points plugins."""
|
||||
self._status = inspect_channels(self.config)
|
||||
enabled_names = {
|
||||
name
|
||||
for name, status in self._status.items()
|
||||
if status.get("enabled") is True
|
||||
}
|
||||
plugin_classes = discover_plugins(enabled_names)
|
||||
builtins = set(discover_channel_names())
|
||||
|
||||
for name in sorted(enabled_names):
|
||||
section = self._get_channel_config(name)
|
||||
if section is None:
|
||||
continue
|
||||
if self._status.get(name, {}).get("available") is False:
|
||||
logger.warning(
|
||||
"%s channel not available: %s",
|
||||
name,
|
||||
self._status[name].get("error", "unavailable"),
|
||||
)
|
||||
continue
|
||||
try:
|
||||
cls = load_channel_class(name) if name in builtins else plugin_classes[name]
|
||||
except Exception as exc: # noqa: BLE001 - status endpoint must explain startup gaps
|
||||
self._status.setdefault(name, {})
|
||||
self._status[name].update(
|
||||
{
|
||||
"available": False,
|
||||
"loaded": False,
|
||||
"running": False,
|
||||
"error": f"{type(exc).__name__}: {exc}",
|
||||
}
|
||||
)
|
||||
logger.warning("%s channel not available: %s", name, exc)
|
||||
continue
|
||||
|
||||
try:
|
||||
kwargs = self._build_channel_kwargs(name)
|
||||
channel = cls(section, self.bus, **kwargs)
|
||||
# Resolve global → per-channel boolean overrides
|
||||
channel.send_progress = self._resolve_bool_override(
|
||||
section, "send_progress", self._global_bool("send_progress", True),
|
||||
)
|
||||
channel.send_tool_hints = self._resolve_bool_override(
|
||||
section, "send_tool_hints", self._global_bool("send_tool_hints", False),
|
||||
)
|
||||
channel.show_reasoning = self._resolve_bool_override(
|
||||
section, "show_reasoning", self._global_bool("show_reasoning", True),
|
||||
)
|
||||
self.channels[name] = channel
|
||||
self._status.setdefault(name, {})
|
||||
self._status[name].update(
|
||||
{
|
||||
"available": True,
|
||||
"loaded": True,
|
||||
"running": channel.is_running,
|
||||
"display_name": getattr(cls, "display_name", name),
|
||||
"error": "",
|
||||
}
|
||||
)
|
||||
logger.info("%s channel enabled", cls.display_name)
|
||||
except Exception as exc:
|
||||
self._status.setdefault(name, {})
|
||||
self._status[name].update(
|
||||
{
|
||||
"loaded": False,
|
||||
"running": False,
|
||||
"error": f"{type(exc).__name__}: {exc}",
|
||||
}
|
||||
)
|
||||
logger.warning("%s channel not available", name, exc_info=True)
|
||||
|
||||
self._validate_allow_from()
|
||||
|
||||
def _build_channel_kwargs(self, name: str) -> dict[str, Any]:
|
||||
"""Build adapter-specific constructor kwargs for channels that need services."""
|
||||
if name == "websocket":
|
||||
from src.channelsui.gateway_services import build_gateway_services
|
||||
|
||||
return {
|
||||
"gateway": build_gateway_services(
|
||||
session_manager=self._session_service,
|
||||
cron_service=self._cron_service,
|
||||
workspace_path=get_workspace_path(),
|
||||
)
|
||||
}
|
||||
if name == "matrix":
|
||||
return {
|
||||
"restrict_to_workspace": self._global_bool("restrict_to_workspace", False),
|
||||
"workspace": get_workspace_path(),
|
||||
}
|
||||
return {}
|
||||
|
||||
def _get_channel_config(self, name: str) -> dict | None:
|
||||
"""Get the config section for channel *name*."""
|
||||
if isinstance(self.config, dict):
|
||||
return self.config.get(name)
|
||||
return getattr(self.config, name, None)
|
||||
|
||||
def _global_bool(self, key: str, default: bool) -> bool:
|
||||
"""Read a global boolean from the top-level channels config."""
|
||||
if isinstance(self.config, dict):
|
||||
return self.config.get(key, default)
|
||||
return getattr(self.config, key, default)
|
||||
|
||||
def _validate_allow_from(self) -> None:
|
||||
for name, ch in self.channels.items():
|
||||
cfg = ch.config
|
||||
if isinstance(cfg, dict):
|
||||
allow = cfg.get("allow_from") if "allow_from" in cfg else cfg.get("allowFrom")
|
||||
else:
|
||||
allow = getattr(cfg, "allow_from", None)
|
||||
if allow is None:
|
||||
logger.info(
|
||||
'"%s" has no allowFrom; unapproved users will receive a pairing code',
|
||||
name,
|
||||
)
|
||||
|
||||
def _should_send_progress(self, channel_name: str, *, tool_hint: bool = False) -> bool:
|
||||
"""Return whether progress (or tool-hints) may be sent to *channel_name*."""
|
||||
ch = self.channels.get(channel_name)
|
||||
if ch is None:
|
||||
logger.debug("Progress check for unknown channel: %s", channel_name)
|
||||
return False
|
||||
return ch.send_tool_hints if tool_hint else ch.send_progress
|
||||
|
||||
def _resolve_bool_override(self, section: Any, key: str, default: bool) -> bool:
|
||||
"""Return *key* from *section* if it is a bool, otherwise *default*.
|
||||
|
||||
For dict configs also checks the camelCase alias.
|
||||
"""
|
||||
if isinstance(section, dict):
|
||||
value = section.get(key)
|
||||
if value is None:
|
||||
camel = _BOOL_CAMEL_ALIASES.get(key)
|
||||
if camel:
|
||||
value = section.get(camel)
|
||||
return value if isinstance(value, bool) else default
|
||||
value = getattr(section, key, None)
|
||||
return value if isinstance(value, bool) else default
|
||||
|
||||
async def _start_channel(self, name: str, channel: BaseChannel) -> None:
|
||||
"""Start a channel and log any exceptions."""
|
||||
try:
|
||||
await channel.start()
|
||||
except Exception:
|
||||
logger.exception("Failed to start channel %s", name)
|
||||
|
||||
async def start_all(self) -> None:
|
||||
"""Start all channels and the outbound dispatcher."""
|
||||
if not self.channels:
|
||||
logger.warning("No channels enabled")
|
||||
return
|
||||
|
||||
self._dispatch_task = asyncio.create_task(self._dispatch_outbound())
|
||||
|
||||
tasks = []
|
||||
for name, channel in self.channels.items():
|
||||
logger.info("Starting %s channel...", name)
|
||||
tasks.append(asyncio.create_task(self._start_channel(name, channel)))
|
||||
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for name, channel in self.channels.items():
|
||||
self._status.setdefault(name, {})
|
||||
self._status[name]["running"] = channel.is_running
|
||||
|
||||
async def stop_all(self) -> None:
|
||||
"""Stop all channels and the dispatcher."""
|
||||
logger.info("Stopping all channels...")
|
||||
|
||||
if self._dispatch_task:
|
||||
self._dispatch_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._dispatch_task
|
||||
|
||||
for name, channel in self.channels.items():
|
||||
try:
|
||||
await channel.stop()
|
||||
logger.info("Stopped %s channel", name)
|
||||
except asyncio.CancelledError:
|
||||
if asyncio.current_task() and asyncio.current_task().cancelling():
|
||||
raise
|
||||
logger.debug("Channel %s stop task was already cancelled", name)
|
||||
except Exception:
|
||||
logger.exception("Error stopping %s", name)
|
||||
finally:
|
||||
self._status.setdefault(name, {})
|
||||
self._status[name]["running"] = channel.is_running
|
||||
|
||||
# --- Outbound dispatch ---
|
||||
|
||||
@staticmethod
|
||||
def _fingerprint_content(content: str) -> str:
|
||||
normalized = " ".join(content.split())
|
||||
return hashlib.sha1(normalized.encode("utf-8")).hexdigest() if normalized else ""
|
||||
|
||||
def _should_suppress_outbound(self, msg: OutboundMessage) -> bool:
|
||||
metadata = msg.metadata or {}
|
||||
if metadata.get("_progress"):
|
||||
return False
|
||||
fingerprint = self._fingerprint_content(msg.content)
|
||||
if not fingerprint:
|
||||
return False
|
||||
|
||||
origin_message_id = metadata.get("origin_message_id")
|
||||
if isinstance(origin_message_id, str) and origin_message_id:
|
||||
key = (msg.channel, msg.chat_id, origin_message_id)
|
||||
if self._origin_reply_fingerprints.get(key) == fingerprint:
|
||||
return True
|
||||
self._origin_reply_fingerprints[key] = fingerprint
|
||||
|
||||
message_id = metadata.get("message_id")
|
||||
if isinstance(message_id, str) and message_id:
|
||||
key = (msg.channel, msg.chat_id, message_id)
|
||||
self._origin_reply_fingerprints[key] = fingerprint
|
||||
|
||||
return False
|
||||
|
||||
async def _dispatch_outbound(self) -> None:
|
||||
"""Dispatch outbound messages to the appropriate channel."""
|
||||
logger.info("Outbound dispatcher started")
|
||||
|
||||
pending: list[OutboundMessage] = []
|
||||
|
||||
while True:
|
||||
try:
|
||||
if pending:
|
||||
msg = pending.pop(0)
|
||||
else:
|
||||
msg = await asyncio.wait_for(
|
||||
self.bus.consume_outbound(), timeout=1.0,
|
||||
)
|
||||
|
||||
# Reasoning routing
|
||||
if (
|
||||
msg.metadata.get("_reasoning_delta")
|
||||
or msg.metadata.get("_reasoning_end")
|
||||
or msg.metadata.get("_reasoning")
|
||||
):
|
||||
channel = self.channels.get(msg.channel)
|
||||
if channel is not None and channel.show_reasoning:
|
||||
await self._send_with_retry(channel, msg)
|
||||
continue
|
||||
|
||||
# Progress / tool-hint filtering
|
||||
if msg.metadata.get("_progress"):
|
||||
if msg.metadata.get("_tool_hint") and not self._should_send_progress(
|
||||
msg.channel, tool_hint=True,
|
||||
):
|
||||
continue
|
||||
if not msg.metadata.get("_tool_hint") and not self._should_send_progress(
|
||||
msg.channel, tool_hint=False,
|
||||
):
|
||||
continue
|
||||
|
||||
if msg.metadata.get("_retry_wait"):
|
||||
continue
|
||||
|
||||
# Coalesce consecutive _stream_delta messages
|
||||
if msg.metadata.get("_stream_delta") and not msg.metadata.get("_stream_end"):
|
||||
msg, extra_pending = self._coalesce_stream_deltas(msg)
|
||||
pending.extend(extra_pending)
|
||||
|
||||
channel = self.channels.get(msg.channel)
|
||||
if channel:
|
||||
if (
|
||||
not msg.metadata.get("_stream_delta")
|
||||
and not msg.metadata.get("_stream_end")
|
||||
and not msg.metadata.get("_streamed")
|
||||
):
|
||||
if self._should_suppress_outbound(msg):
|
||||
logger.info(
|
||||
"Suppressing duplicate outbound message to %s:%s",
|
||||
msg.channel, msg.chat_id,
|
||||
)
|
||||
continue
|
||||
await self._send_with_retry(channel, msg)
|
||||
else:
|
||||
logger.warning("Unknown channel: %s", msg.channel)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
|
||||
@staticmethod
|
||||
async def _send_once(channel: BaseChannel, msg: OutboundMessage) -> None:
|
||||
"""Send one outbound message without retry policy."""
|
||||
if msg.metadata.get("_reasoning_end"):
|
||||
await channel.send_reasoning_end(msg.chat_id, msg.metadata)
|
||||
elif msg.metadata.get("_reasoning_delta"):
|
||||
await channel.send_reasoning_delta(msg.chat_id, msg.content, msg.metadata)
|
||||
elif msg.metadata.get("_reasoning"):
|
||||
await channel.send_reasoning(msg)
|
||||
elif msg.metadata.get("_file_edit_events"):
|
||||
edits = msg.metadata.get("_file_edit_events")
|
||||
await channel.send_file_edit_events(
|
||||
msg.chat_id,
|
||||
edits if isinstance(edits, list) else [],
|
||||
msg.metadata,
|
||||
)
|
||||
elif msg.metadata.get("_stream_delta") or msg.metadata.get("_stream_end"):
|
||||
await channel.send_delta(msg.chat_id, msg.content, msg.metadata)
|
||||
elif not msg.metadata.get("_streamed"):
|
||||
await channel.send(msg)
|
||||
|
||||
def _coalesce_stream_deltas(
|
||||
self, first_msg: OutboundMessage,
|
||||
) -> tuple[OutboundMessage, list[OutboundMessage]]:
|
||||
"""Merge consecutive _stream_delta messages for the same target.
|
||||
|
||||
This reduces the number of API calls when the queue has accumulated
|
||||
multiple deltas.
|
||||
"""
|
||||
first_metadata = first_msg.metadata or {}
|
||||
target_key = (
|
||||
first_msg.channel,
|
||||
first_msg.chat_id,
|
||||
first_metadata.get("_stream_id"),
|
||||
)
|
||||
combined_content = first_msg.content
|
||||
final_metadata = dict(first_msg.metadata or {})
|
||||
non_matching: list[OutboundMessage] = []
|
||||
|
||||
while True:
|
||||
try:
|
||||
next_msg = self.bus.outbound.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
next_metadata = next_msg.metadata or {}
|
||||
same_target = (
|
||||
next_msg.channel,
|
||||
next_msg.chat_id,
|
||||
next_metadata.get("_stream_id"),
|
||||
) == target_key
|
||||
is_delta = next_metadata.get("_stream_delta")
|
||||
is_end = next_metadata.get("_stream_end")
|
||||
|
||||
if same_target and is_delta and not final_metadata.get("_stream_end"):
|
||||
combined_content += next_msg.content
|
||||
if is_end:
|
||||
final_metadata["_stream_end"] = True
|
||||
break
|
||||
else:
|
||||
non_matching.append(next_msg)
|
||||
break
|
||||
|
||||
merged = OutboundMessage(
|
||||
channel=first_msg.channel,
|
||||
chat_id=first_msg.chat_id,
|
||||
content=combined_content,
|
||||
metadata=final_metadata,
|
||||
)
|
||||
return merged, non_matching
|
||||
|
||||
async def _send_with_retry(
|
||||
self, channel: BaseChannel, msg: OutboundMessage,
|
||||
) -> None:
|
||||
"""Send a message with retry on failure using exponential backoff."""
|
||||
max_attempts = 2 # conservative default
|
||||
if isinstance(self.config, dict):
|
||||
max_attempts = max(self.config.get("send_max_retries", 2), 1)
|
||||
elif hasattr(self.config, "send_max_retries"):
|
||||
max_attempts = max(self.config.send_max_retries, 1)
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
try:
|
||||
await self._send_once(channel, msg)
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
if attempt == max_attempts - 1:
|
||||
logger.exception(
|
||||
"Failed to send to %s after %d attempts",
|
||||
msg.channel, max_attempts,
|
||||
)
|
||||
return
|
||||
delay = _SEND_RETRY_DELAYS[min(attempt, len(_SEND_RETRY_DELAYS) - 1)]
|
||||
logger.warning(
|
||||
"Send to %s failed (attempt %d/%d): %s, retrying in %ds",
|
||||
msg.channel, attempt + 1, max_attempts, type(e).__name__, delay,
|
||||
)
|
||||
try:
|
||||
await asyncio.sleep(delay)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
|
||||
# --- Public API ---
|
||||
|
||||
def get_channel(self, name: str) -> BaseChannel | None:
|
||||
"""Get a channel by name."""
|
||||
return self.channels.get(name)
|
||||
|
||||
def get_status(self) -> dict[str, Any]:
|
||||
"""Get status of all channels."""
|
||||
status = {name: dict(item) for name, item in self._status.items()}
|
||||
for name, channel in self.channels.items():
|
||||
status.setdefault(name, {})
|
||||
status[name].update(
|
||||
{
|
||||
"enabled": True,
|
||||
"loaded": True,
|
||||
"running": channel.is_running,
|
||||
"display_name": getattr(channel, "display_name", name),
|
||||
}
|
||||
)
|
||||
return status
|
||||
|
||||
@property
|
||||
def enabled_channels(self) -> list[str]:
|
||||
"""Get list of enabled channel names."""
|
||||
return list(self.channels.keys())
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,943 @@
|
||||
"""Mochat channel implementation using Socket.IO with HTTP polling fallback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections import deque
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_runtime_subdir
|
||||
from pydantic import BaseModel
|
||||
|
||||
try:
|
||||
import socketio
|
||||
SOCKETIO_AVAILABLE = True
|
||||
except ImportError:
|
||||
socketio = None
|
||||
SOCKETIO_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import msgpack # noqa: F401
|
||||
MSGPACK_AVAILABLE = True
|
||||
except ImportError:
|
||||
MSGPACK_AVAILABLE = False
|
||||
|
||||
MAX_SEEN_MESSAGE_IDS = 2000
|
||||
CURSOR_SAVE_DEBOUNCE_S = 0.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data classes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class MochatBufferedEntry:
|
||||
"""Buffered inbound entry for delayed dispatch."""
|
||||
raw_body: str
|
||||
author: str
|
||||
sender_name: str = ""
|
||||
sender_username: str = ""
|
||||
timestamp: int | None = None
|
||||
message_id: str = ""
|
||||
group_id: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class DelayState:
|
||||
"""Per-target delayed message state."""
|
||||
entries: list[MochatBufferedEntry] = field(default_factory=list)
|
||||
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
timer: asyncio.Task | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MochatTarget:
|
||||
"""Outbound target resolution result."""
|
||||
id: str
|
||||
is_panel: bool
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _safe_dict(value: Any) -> dict:
|
||||
"""Return *value* if it's a dict, else empty dict."""
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
|
||||
def _str_field(src: dict, *keys: str) -> str:
|
||||
"""Return the first non-empty str value found for *keys*, stripped."""
|
||||
for k in keys:
|
||||
v = src.get(k)
|
||||
if isinstance(v, str) and v.strip():
|
||||
return v.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _make_synthetic_event(
|
||||
message_id: str, author: str, content: Any,
|
||||
meta: Any, group_id: str, converse_id: str,
|
||||
timestamp: Any = None, *, author_info: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a synthetic ``message.add`` event dict."""
|
||||
payload: dict[str, Any] = {
|
||||
"messageId": message_id, "author": author,
|
||||
"content": content, "meta": _safe_dict(meta),
|
||||
"groupId": group_id, "converseId": converse_id,
|
||||
}
|
||||
if author_info is not None:
|
||||
payload["authorInfo"] = _safe_dict(author_info)
|
||||
return {
|
||||
"type": "message.add",
|
||||
"timestamp": timestamp or datetime.utcnow().isoformat(),
|
||||
"payload": payload,
|
||||
}
|
||||
|
||||
|
||||
def normalize_mochat_content(content: Any) -> str:
|
||||
"""Normalize content payload to text."""
|
||||
if isinstance(content, str):
|
||||
return content.strip()
|
||||
if content is None:
|
||||
return ""
|
||||
try:
|
||||
return json.dumps(content, ensure_ascii=False)
|
||||
except TypeError:
|
||||
return str(content)
|
||||
|
||||
|
||||
def resolve_mochat_target(raw: str) -> MochatTarget:
|
||||
"""Resolve id and target kind from user-provided target string."""
|
||||
trimmed = (raw or "").strip()
|
||||
if not trimmed:
|
||||
return MochatTarget(id="", is_panel=False)
|
||||
|
||||
lowered = trimmed.lower()
|
||||
cleaned, forced_panel = trimmed, False
|
||||
for prefix in ("mochat:", "group:", "channel:", "panel:"):
|
||||
if lowered.startswith(prefix):
|
||||
cleaned = trimmed[len(prefix):].strip()
|
||||
forced_panel = prefix in {"group:", "channel:", "panel:"}
|
||||
break
|
||||
|
||||
if not cleaned:
|
||||
return MochatTarget(id="", is_panel=False)
|
||||
return MochatTarget(id=cleaned, is_panel=forced_panel or not cleaned.startswith("session_"))
|
||||
|
||||
|
||||
def extract_mention_ids(value: Any) -> list[str]:
|
||||
"""Extract mention ids from heterogeneous mention payload."""
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
ids: list[str] = []
|
||||
for item in value:
|
||||
if isinstance(item, str):
|
||||
if item.strip():
|
||||
ids.append(item.strip())
|
||||
elif isinstance(item, dict):
|
||||
for key in ("id", "userId", "_id"):
|
||||
candidate = item.get(key)
|
||||
if isinstance(candidate, str) and candidate.strip():
|
||||
ids.append(candidate.strip())
|
||||
break
|
||||
return ids
|
||||
|
||||
|
||||
def resolve_was_mentioned(payload: dict[str, Any], agent_user_id: str) -> bool:
|
||||
"""Resolve mention state from payload metadata and text fallback."""
|
||||
meta = payload.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
if meta.get("mentioned") is True or meta.get("wasMentioned") is True:
|
||||
return True
|
||||
for f in ("mentions", "mentionIds", "mentionedUserIds", "mentionedUsers"):
|
||||
if agent_user_id and agent_user_id in extract_mention_ids(meta.get(f)):
|
||||
return True
|
||||
if not agent_user_id:
|
||||
return False
|
||||
content = payload.get("content")
|
||||
if not isinstance(content, str) or not content:
|
||||
return False
|
||||
return f"<@{agent_user_id}>" in content or f"@{agent_user_id}" in content
|
||||
|
||||
|
||||
def resolve_require_mention(config: MochatConfig, session_id: str, group_id: str) -> bool:
|
||||
"""Resolve mention requirement for group/panel conversations."""
|
||||
groups = config.groups or {}
|
||||
for key in (group_id, session_id, "*"):
|
||||
if key and key in groups:
|
||||
return bool(groups[key].require_mention)
|
||||
return bool(config.mention.require_in_groups)
|
||||
|
||||
|
||||
def build_buffered_body(entries: list[MochatBufferedEntry], is_group: bool) -> str:
|
||||
"""Build text body from one or more buffered entries."""
|
||||
if not entries:
|
||||
return ""
|
||||
if len(entries) == 1:
|
||||
return entries[0].raw_body
|
||||
lines: list[str] = []
|
||||
for entry in entries:
|
||||
if not entry.raw_body:
|
||||
continue
|
||||
if is_group:
|
||||
label = entry.sender_name.strip() or entry.sender_username.strip() or entry.author
|
||||
if label:
|
||||
lines.append(f"{label}: {entry.raw_body}")
|
||||
continue
|
||||
lines.append(entry.raw_body)
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def parse_timestamp(value: Any) -> int | None:
|
||||
"""Parse event timestamp to epoch milliseconds."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
try:
|
||||
return int(datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() * 1000)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config classes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MochatMentionConfig(BaseModel):
|
||||
"""Mochat mention behavior configuration."""
|
||||
|
||||
require_in_groups: bool = False
|
||||
|
||||
|
||||
class MochatGroupRule(BaseModel):
|
||||
"""Mochat per-group mention requirement."""
|
||||
|
||||
require_mention: bool = False
|
||||
|
||||
|
||||
class MochatConfig(BaseModel):
|
||||
"""Mochat channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
base_url: str = "https://mochat.io"
|
||||
socket_url: str = ""
|
||||
socket_path: str = "/socket.io"
|
||||
socket_disable_msgpack: bool = False
|
||||
socket_reconnect_delay_ms: int = 1000
|
||||
socket_max_reconnect_delay_ms: int = 10000
|
||||
socket_connect_timeout_ms: int = 10000
|
||||
refresh_interval_ms: int = 30000
|
||||
watch_timeout_ms: int = 25000
|
||||
watch_limit: int = 100
|
||||
retry_delay_ms: int = 500
|
||||
max_retry_attempts: int = 0
|
||||
claw_token: str = ""
|
||||
agent_user_id: str = ""
|
||||
sessions: list[str] = Field(default_factory=list)
|
||||
panels: list[str] = Field(default_factory=list)
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
mention: MochatMentionConfig = Field(default_factory=MochatMentionConfig)
|
||||
groups: dict[str, MochatGroupRule] = Field(default_factory=dict)
|
||||
reply_delay_mode: str = "non-mention"
|
||||
reply_delay_ms: int = 120000
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Channel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MochatChannel(BaseChannel):
|
||||
"""Mochat channel using socket.io with fallback polling workers."""
|
||||
|
||||
name = "mochat"
|
||||
display_name = "Mochat"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return MochatConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = MochatConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: MochatConfig = config
|
||||
self._http: httpx.AsyncClient | None = None
|
||||
self._socket: Any = None
|
||||
self._ws_connected = self._ws_ready = False
|
||||
|
||||
self._state_dir = get_runtime_subdir("mochat")
|
||||
self._cursor_path = self._state_dir / "session_cursors.json"
|
||||
self._session_cursor: dict[str, int] = {}
|
||||
self._cursor_save_task: asyncio.Task | None = None
|
||||
|
||||
self._session_set: set[str] = set()
|
||||
self._panel_set: set[str] = set()
|
||||
self._auto_discover_sessions = self._auto_discover_panels = False
|
||||
|
||||
self._cold_sessions: set[str] = set()
|
||||
self._session_by_converse: dict[str, str] = {}
|
||||
|
||||
self._seen_set: dict[str, set[str]] = {}
|
||||
self._seen_queue: dict[str, deque[str]] = {}
|
||||
self._delay_states: dict[str, DelayState] = {}
|
||||
|
||||
self._fallback_mode = False
|
||||
self._session_fallback_tasks: dict[str, asyncio.Task] = {}
|
||||
self._panel_fallback_tasks: dict[str, asyncio.Task] = {}
|
||||
self._refresh_task: asyncio.Task | None = None
|
||||
self._target_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
# ---- lifecycle ---------------------------------------------------------
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start Mochat channel workers and websocket connection."""
|
||||
if not self.config.claw_token:
|
||||
self.logger.error("claw_token not configured")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._http = httpx.AsyncClient(timeout=30.0)
|
||||
self._state_dir.mkdir(parents=True, exist_ok=True)
|
||||
await self._load_session_cursors()
|
||||
self._seed_targets_from_config()
|
||||
await self._refresh_targets(subscribe_new=False)
|
||||
|
||||
if not await self._start_socket_client():
|
||||
await self._ensure_fallback_workers()
|
||||
|
||||
self._refresh_task = asyncio.create_task(self._refresh_loop())
|
||||
while self._running:
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop all workers and clean up resources."""
|
||||
self._running = False
|
||||
if self._refresh_task:
|
||||
self._refresh_task.cancel()
|
||||
self._refresh_task = None
|
||||
|
||||
await self._stop_fallback_workers()
|
||||
await self._cancel_delay_timers()
|
||||
|
||||
if self._socket:
|
||||
with suppress(Exception):
|
||||
await self._socket.disconnect()
|
||||
self._socket = None
|
||||
|
||||
if self._cursor_save_task:
|
||||
self._cursor_save_task.cancel()
|
||||
self._cursor_save_task = None
|
||||
await self._save_session_cursors()
|
||||
|
||||
if self._http:
|
||||
await self._http.aclose()
|
||||
self._http = None
|
||||
self._ws_connected = self._ws_ready = False
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send outbound message to session or panel."""
|
||||
if not self.config.claw_token:
|
||||
self.logger.warning("claw_token missing, skip send")
|
||||
return
|
||||
|
||||
parts = ([msg.content.strip()] if msg.content and msg.content.strip() else [])
|
||||
if msg.media:
|
||||
parts.extend(m for m in msg.media if isinstance(m, str) and m.strip())
|
||||
content = "\n".join(parts).strip()
|
||||
if not content:
|
||||
return
|
||||
|
||||
target = resolve_mochat_target(msg.chat_id)
|
||||
if not target.id:
|
||||
self.logger.warning("outbound target is empty")
|
||||
return
|
||||
|
||||
is_panel = (target.is_panel or target.id in self._panel_set) and not target.id.startswith("session_")
|
||||
try:
|
||||
if is_panel:
|
||||
await self._api_send("/api/claw/groups/panels/send", "panelId", target.id,
|
||||
content, msg.reply_to, self._read_group_id(msg.metadata))
|
||||
else:
|
||||
await self._api_send("/api/claw/sessions/send", "sessionId", target.id,
|
||||
content, msg.reply_to)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to send message")
|
||||
raise
|
||||
|
||||
# ---- config / init helpers ---------------------------------------------
|
||||
|
||||
def _seed_targets_from_config(self) -> None:
|
||||
sessions, self._auto_discover_sessions = self._normalize_id_list(self.config.sessions)
|
||||
panels, self._auto_discover_panels = self._normalize_id_list(self.config.panels)
|
||||
self._session_set.update(sessions)
|
||||
self._panel_set.update(panels)
|
||||
for sid in sessions:
|
||||
if sid not in self._session_cursor:
|
||||
self._cold_sessions.add(sid)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_id_list(values: list[str]) -> tuple[list[str], bool]:
|
||||
cleaned = [str(v).strip() for v in values if str(v).strip()]
|
||||
return sorted({v for v in cleaned if v != "*"}), "*" in cleaned
|
||||
|
||||
# ---- websocket ---------------------------------------------------------
|
||||
|
||||
async def _start_socket_client(self) -> bool:
|
||||
if not SOCKETIO_AVAILABLE:
|
||||
self.logger.warning("python-socketio not installed, using polling fallback")
|
||||
return False
|
||||
|
||||
serializer = "default"
|
||||
if not self.config.socket_disable_msgpack:
|
||||
if MSGPACK_AVAILABLE:
|
||||
serializer = "msgpack"
|
||||
else:
|
||||
self.logger.warning("msgpack not installed but socket_disable_msgpack=false; using JSON")
|
||||
|
||||
client = socketio.AsyncClient(
|
||||
reconnection=True,
|
||||
reconnection_attempts=self.config.max_retry_attempts or None,
|
||||
reconnection_delay=max(0.1, self.config.socket_reconnect_delay_ms / 1000.0),
|
||||
reconnection_delay_max=max(0.1, self.config.socket_max_reconnect_delay_ms / 1000.0),
|
||||
logger=False, engineio_logger=False, serializer=serializer,
|
||||
)
|
||||
|
||||
@client.event
|
||||
async def connect() -> None:
|
||||
self._ws_connected, self._ws_ready = True, False
|
||||
self.logger.info("websocket connected")
|
||||
subscribed = await self._subscribe_all()
|
||||
self._ws_ready = subscribed
|
||||
await (self._stop_fallback_workers() if subscribed else self._ensure_fallback_workers())
|
||||
|
||||
@client.event
|
||||
async def disconnect() -> None:
|
||||
if not self._running:
|
||||
return
|
||||
self._ws_connected = self._ws_ready = False
|
||||
self.logger.warning("websocket disconnected")
|
||||
await self._ensure_fallback_workers()
|
||||
|
||||
@client.event
|
||||
async def connect_error(data: Any) -> None:
|
||||
self.logger.error("websocket connect error: {}", data)
|
||||
|
||||
@client.on("claw.session.events")
|
||||
async def on_session_events(payload: dict[str, Any]) -> None:
|
||||
await self._handle_watch_payload(payload, "session")
|
||||
|
||||
@client.on("claw.panel.events")
|
||||
async def on_panel_events(payload: dict[str, Any]) -> None:
|
||||
await self._handle_watch_payload(payload, "panel")
|
||||
|
||||
for ev in ("notify:chat.inbox.append", "notify:chat.message.add",
|
||||
"notify:chat.message.update", "notify:chat.message.recall",
|
||||
"notify:chat.message.delete"):
|
||||
client.on(ev, self._build_notify_handler(ev))
|
||||
|
||||
socket_url = (self.config.socket_url or self.config.base_url).strip().rstrip("/")
|
||||
socket_path = (self.config.socket_path or "/socket.io").strip().lstrip("/")
|
||||
|
||||
try:
|
||||
self._socket = client
|
||||
await client.connect(
|
||||
socket_url, transports=["websocket"], socketio_path=socket_path,
|
||||
auth={"token": self.config.claw_token},
|
||||
wait_timeout=max(1.0, self.config.socket_connect_timeout_ms / 1000.0),
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
self.logger.exception("Failed to connect websocket")
|
||||
with suppress(Exception):
|
||||
await client.disconnect()
|
||||
self._socket = None
|
||||
return False
|
||||
|
||||
def _build_notify_handler(self, event_name: str):
|
||||
async def handler(payload: Any) -> None:
|
||||
if event_name == "notify:chat.inbox.append":
|
||||
await self._handle_notify_inbox_append(payload)
|
||||
elif event_name.startswith("notify:chat.message."):
|
||||
await self._handle_notify_chat_message(payload)
|
||||
return handler
|
||||
|
||||
# ---- subscribe ---------------------------------------------------------
|
||||
|
||||
async def _subscribe_all(self) -> bool:
|
||||
ok = await self._subscribe_sessions(sorted(self._session_set))
|
||||
ok = await self._subscribe_panels(sorted(self._panel_set)) and ok
|
||||
if self._auto_discover_sessions or self._auto_discover_panels:
|
||||
await self._refresh_targets(subscribe_new=True)
|
||||
return ok
|
||||
|
||||
async def _subscribe_sessions(self, session_ids: list[str]) -> bool:
|
||||
if not session_ids:
|
||||
return True
|
||||
for sid in session_ids:
|
||||
if sid not in self._session_cursor:
|
||||
self._cold_sessions.add(sid)
|
||||
|
||||
ack = await self._socket_call("com.claw.im.subscribeSessions", {
|
||||
"sessionIds": session_ids, "cursors": self._session_cursor,
|
||||
"limit": self.config.watch_limit,
|
||||
})
|
||||
if not ack.get("result"):
|
||||
self.logger.error("subscribeSessions failed: {}", ack.get('message', 'unknown error'))
|
||||
return False
|
||||
|
||||
data = ack.get("data")
|
||||
items: list[dict[str, Any]] = []
|
||||
if isinstance(data, list):
|
||||
items = [i for i in data if isinstance(i, dict)]
|
||||
elif isinstance(data, dict):
|
||||
sessions = data.get("sessions")
|
||||
if isinstance(sessions, list):
|
||||
items = [i for i in sessions if isinstance(i, dict)]
|
||||
elif "sessionId" in data:
|
||||
items = [data]
|
||||
for p in items:
|
||||
await self._handle_watch_payload(p, "session")
|
||||
return True
|
||||
|
||||
async def _subscribe_panels(self, panel_ids: list[str]) -> bool:
|
||||
if not self._auto_discover_panels and not panel_ids:
|
||||
return True
|
||||
ack = await self._socket_call("com.claw.im.subscribePanels", {"panelIds": panel_ids})
|
||||
if not ack.get("result"):
|
||||
self.logger.error("subscribePanels failed: {}", ack.get('message', 'unknown error'))
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _socket_call(self, event_name: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
if not self._socket:
|
||||
return {"result": False, "message": "socket not connected"}
|
||||
try:
|
||||
raw = await self._socket.call(event_name, payload, timeout=10)
|
||||
except Exception as e:
|
||||
return {"result": False, "message": str(e)}
|
||||
return raw if isinstance(raw, dict) else {"result": True, "data": raw}
|
||||
|
||||
# ---- refresh / discovery -----------------------------------------------
|
||||
|
||||
async def _refresh_loop(self) -> None:
|
||||
interval_s = max(1.0, self.config.refresh_interval_ms / 1000.0)
|
||||
while self._running:
|
||||
await asyncio.sleep(interval_s)
|
||||
try:
|
||||
await self._refresh_targets(subscribe_new=self._ws_ready)
|
||||
except Exception as e:
|
||||
self.logger.warning("refresh failed: {}", e)
|
||||
if self._fallback_mode:
|
||||
await self._ensure_fallback_workers()
|
||||
|
||||
async def _refresh_targets(self, subscribe_new: bool) -> None:
|
||||
if self._auto_discover_sessions:
|
||||
await self._refresh_sessions_directory(subscribe_new)
|
||||
if self._auto_discover_panels:
|
||||
await self._refresh_panels(subscribe_new)
|
||||
|
||||
async def _refresh_sessions_directory(self, subscribe_new: bool) -> None:
|
||||
try:
|
||||
response = await self._post_json("/api/claw/sessions/list", {})
|
||||
except Exception as e:
|
||||
self.logger.warning("listSessions failed: {}", e)
|
||||
return
|
||||
|
||||
sessions = response.get("sessions")
|
||||
if not isinstance(sessions, list):
|
||||
return
|
||||
|
||||
new_ids: list[str] = []
|
||||
for s in sessions:
|
||||
if not isinstance(s, dict):
|
||||
continue
|
||||
sid = _str_field(s, "sessionId")
|
||||
if not sid:
|
||||
continue
|
||||
if sid not in self._session_set:
|
||||
self._session_set.add(sid)
|
||||
new_ids.append(sid)
|
||||
if sid not in self._session_cursor:
|
||||
self._cold_sessions.add(sid)
|
||||
cid = _str_field(s, "converseId")
|
||||
if cid:
|
||||
self._session_by_converse[cid] = sid
|
||||
|
||||
if not new_ids:
|
||||
return
|
||||
if self._ws_ready and subscribe_new:
|
||||
await self._subscribe_sessions(new_ids)
|
||||
if self._fallback_mode:
|
||||
await self._ensure_fallback_workers()
|
||||
|
||||
async def _refresh_panels(self, subscribe_new: bool) -> None:
|
||||
try:
|
||||
response = await self._post_json("/api/claw/groups/get", {})
|
||||
except Exception as e:
|
||||
self.logger.warning("getWorkspaceGroup failed: {}", e)
|
||||
return
|
||||
|
||||
raw_panels = response.get("panels")
|
||||
if not isinstance(raw_panels, list):
|
||||
return
|
||||
|
||||
new_ids: list[str] = []
|
||||
for p in raw_panels:
|
||||
if not isinstance(p, dict):
|
||||
continue
|
||||
pt = p.get("type")
|
||||
if isinstance(pt, int) and pt != 0:
|
||||
continue
|
||||
pid = _str_field(p, "id", "_id")
|
||||
if pid and pid not in self._panel_set:
|
||||
self._panel_set.add(pid)
|
||||
new_ids.append(pid)
|
||||
|
||||
if not new_ids:
|
||||
return
|
||||
if self._ws_ready and subscribe_new:
|
||||
await self._subscribe_panels(new_ids)
|
||||
if self._fallback_mode:
|
||||
await self._ensure_fallback_workers()
|
||||
|
||||
# ---- fallback workers --------------------------------------------------
|
||||
|
||||
async def _ensure_fallback_workers(self) -> None:
|
||||
if not self._running:
|
||||
return
|
||||
self._fallback_mode = True
|
||||
for sid in sorted(self._session_set):
|
||||
t = self._session_fallback_tasks.get(sid)
|
||||
if not t or t.done():
|
||||
self._session_fallback_tasks[sid] = asyncio.create_task(self._session_watch_worker(sid))
|
||||
for pid in sorted(self._panel_set):
|
||||
t = self._panel_fallback_tasks.get(pid)
|
||||
if not t or t.done():
|
||||
self._panel_fallback_tasks[pid] = asyncio.create_task(self._panel_poll_worker(pid))
|
||||
|
||||
async def _stop_fallback_workers(self) -> None:
|
||||
self._fallback_mode = False
|
||||
tasks = [*self._session_fallback_tasks.values(), *self._panel_fallback_tasks.values()]
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._session_fallback_tasks.clear()
|
||||
self._panel_fallback_tasks.clear()
|
||||
|
||||
async def _session_watch_worker(self, session_id: str) -> None:
|
||||
while self._running and self._fallback_mode:
|
||||
try:
|
||||
payload = await self._post_json("/api/claw/sessions/watch", {
|
||||
"sessionId": session_id, "cursor": self._session_cursor.get(session_id, 0),
|
||||
"timeoutMs": self.config.watch_timeout_ms, "limit": self.config.watch_limit,
|
||||
})
|
||||
await self._handle_watch_payload(payload, "session")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
self.logger.warning("watch fallback error ({}): {}", session_id, e)
|
||||
await asyncio.sleep(max(0.1, self.config.retry_delay_ms / 1000.0))
|
||||
|
||||
async def _panel_poll_worker(self, panel_id: str) -> None:
|
||||
sleep_s = max(1.0, self.config.refresh_interval_ms / 1000.0)
|
||||
while self._running and self._fallback_mode:
|
||||
try:
|
||||
resp = await self._post_json("/api/claw/groups/panels/messages", {
|
||||
"panelId": panel_id, "limit": min(100, max(1, self.config.watch_limit)),
|
||||
})
|
||||
msgs = resp.get("messages")
|
||||
if isinstance(msgs, list):
|
||||
for m in reversed(msgs):
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
evt = _make_synthetic_event(
|
||||
message_id=str(m.get("messageId") or ""),
|
||||
author=str(m.get("author") or ""),
|
||||
content=m.get("content"),
|
||||
meta=m.get("meta"), group_id=str(resp.get("groupId") or ""),
|
||||
converse_id=panel_id, timestamp=m.get("createdAt"),
|
||||
author_info=m.get("authorInfo"),
|
||||
)
|
||||
await self._process_inbound_event(panel_id, evt, "panel")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
self.logger.warning("panel polling error ({}): {}", panel_id, e)
|
||||
await asyncio.sleep(sleep_s)
|
||||
|
||||
# ---- inbound event processing ------------------------------------------
|
||||
|
||||
async def _handle_watch_payload(self, payload: dict[str, Any], target_kind: str) -> None:
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
target_id = _str_field(payload, "sessionId")
|
||||
if not target_id:
|
||||
return
|
||||
|
||||
lock = self._target_locks.setdefault(f"{target_kind}:{target_id}", asyncio.Lock())
|
||||
async with lock:
|
||||
prev = self._session_cursor.get(target_id, 0) if target_kind == "session" else 0
|
||||
pc = payload.get("cursor")
|
||||
if target_kind == "session" and isinstance(pc, int) and pc >= 0:
|
||||
self._mark_session_cursor(target_id, pc)
|
||||
|
||||
raw_events = payload.get("events")
|
||||
if not isinstance(raw_events, list):
|
||||
return
|
||||
if target_kind == "session" and target_id in self._cold_sessions:
|
||||
self._cold_sessions.discard(target_id)
|
||||
return
|
||||
|
||||
for event in raw_events:
|
||||
if not isinstance(event, dict):
|
||||
continue
|
||||
seq = event.get("seq")
|
||||
if target_kind == "session" and isinstance(seq, int) and seq > self._session_cursor.get(target_id, prev):
|
||||
self._mark_session_cursor(target_id, seq)
|
||||
if event.get("type") == "message.add":
|
||||
await self._process_inbound_event(target_id, event, target_kind)
|
||||
|
||||
async def _process_inbound_event(self, target_id: str, event: dict[str, Any], target_kind: str) -> None:
|
||||
payload = event.get("payload")
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
|
||||
author = _str_field(payload, "author")
|
||||
if not author or (self.config.agent_user_id and author == self.config.agent_user_id):
|
||||
return
|
||||
if not self.is_allowed(author):
|
||||
return
|
||||
|
||||
message_id = _str_field(payload, "messageId")
|
||||
seen_key = f"{target_kind}:{target_id}"
|
||||
if message_id and self._remember_message_id(seen_key, message_id):
|
||||
return
|
||||
|
||||
raw_body = normalize_mochat_content(payload.get("content")) or "[empty message]"
|
||||
ai = _safe_dict(payload.get("authorInfo"))
|
||||
sender_name = _str_field(ai, "nickname", "email")
|
||||
sender_username = _str_field(ai, "agentId")
|
||||
|
||||
group_id = _str_field(payload, "groupId")
|
||||
is_group = bool(group_id)
|
||||
was_mentioned = resolve_was_mentioned(payload, self.config.agent_user_id)
|
||||
require_mention = target_kind == "panel" and is_group and resolve_require_mention(self.config, target_id, group_id)
|
||||
use_delay = target_kind == "panel" and self.config.reply_delay_mode == "non-mention"
|
||||
|
||||
if require_mention and not was_mentioned and not use_delay:
|
||||
return
|
||||
|
||||
entry = MochatBufferedEntry(
|
||||
raw_body=raw_body, author=author, sender_name=sender_name,
|
||||
sender_username=sender_username, timestamp=parse_timestamp(event.get("timestamp")),
|
||||
message_id=message_id, group_id=group_id,
|
||||
)
|
||||
|
||||
if use_delay:
|
||||
delay_key = seen_key
|
||||
if was_mentioned:
|
||||
await self._flush_delayed_entries(delay_key, target_id, target_kind, "mention", entry)
|
||||
else:
|
||||
await self._enqueue_delayed_entry(delay_key, target_id, target_kind, entry)
|
||||
return
|
||||
|
||||
await self._dispatch_entries(target_id, target_kind, [entry], was_mentioned)
|
||||
|
||||
# ---- dedup / buffering -------------------------------------------------
|
||||
|
||||
def _remember_message_id(self, key: str, message_id: str) -> bool:
|
||||
seen_set = self._seen_set.setdefault(key, set())
|
||||
seen_queue = self._seen_queue.setdefault(key, deque())
|
||||
if message_id in seen_set:
|
||||
return True
|
||||
seen_set.add(message_id)
|
||||
seen_queue.append(message_id)
|
||||
while len(seen_queue) > MAX_SEEN_MESSAGE_IDS:
|
||||
seen_set.discard(seen_queue.popleft())
|
||||
return False
|
||||
|
||||
async def _enqueue_delayed_entry(self, key: str, target_id: str, target_kind: str, entry: MochatBufferedEntry) -> None:
|
||||
state = self._delay_states.setdefault(key, DelayState())
|
||||
async with state.lock:
|
||||
state.entries.append(entry)
|
||||
if state.timer:
|
||||
state.timer.cancel()
|
||||
state.timer = asyncio.create_task(self._delay_flush_after(key, target_id, target_kind))
|
||||
|
||||
async def _delay_flush_after(self, key: str, target_id: str, target_kind: str) -> None:
|
||||
await asyncio.sleep(max(0, self.config.reply_delay_ms) / 1000.0)
|
||||
await self._flush_delayed_entries(key, target_id, target_kind, "timer", None)
|
||||
|
||||
async def _flush_delayed_entries(self, key: str, target_id: str, target_kind: str, reason: str, entry: MochatBufferedEntry | None) -> None:
|
||||
state = self._delay_states.setdefault(key, DelayState())
|
||||
async with state.lock:
|
||||
if entry:
|
||||
state.entries.append(entry)
|
||||
current = asyncio.current_task()
|
||||
if state.timer and state.timer is not current:
|
||||
state.timer.cancel()
|
||||
state.timer = None
|
||||
entries = state.entries[:]
|
||||
state.entries.clear()
|
||||
if entries:
|
||||
await self._dispatch_entries(target_id, target_kind, entries, reason == "mention")
|
||||
|
||||
async def _dispatch_entries(self, target_id: str, target_kind: str, entries: list[MochatBufferedEntry], was_mentioned: bool) -> None:
|
||||
if not entries:
|
||||
return
|
||||
last = entries[-1]
|
||||
is_group = bool(last.group_id)
|
||||
body = build_buffered_body(entries, is_group) or "[empty message]"
|
||||
await self._handle_message(
|
||||
sender_id=last.author, chat_id=target_id, content=body,
|
||||
metadata={
|
||||
"message_id": last.message_id, "timestamp": last.timestamp,
|
||||
"is_group": is_group, "group_id": last.group_id,
|
||||
"sender_name": last.sender_name, "sender_username": last.sender_username,
|
||||
"target_kind": target_kind, "was_mentioned": was_mentioned,
|
||||
"buffered_count": len(entries),
|
||||
},
|
||||
)
|
||||
|
||||
async def _cancel_delay_timers(self) -> None:
|
||||
for state in self._delay_states.values():
|
||||
if state.timer:
|
||||
state.timer.cancel()
|
||||
self._delay_states.clear()
|
||||
|
||||
# ---- notify handlers ---------------------------------------------------
|
||||
|
||||
async def _handle_notify_chat_message(self, payload: Any) -> None:
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
group_id = _str_field(payload, "groupId")
|
||||
panel_id = _str_field(payload, "converseId", "panelId")
|
||||
if not group_id or not panel_id:
|
||||
return
|
||||
if self._panel_set and panel_id not in self._panel_set:
|
||||
return
|
||||
|
||||
evt = _make_synthetic_event(
|
||||
message_id=str(payload.get("_id") or payload.get("messageId") or ""),
|
||||
author=str(payload.get("author") or ""),
|
||||
content=payload.get("content"), meta=payload.get("meta"),
|
||||
group_id=group_id, converse_id=panel_id,
|
||||
timestamp=payload.get("createdAt"), author_info=payload.get("authorInfo"),
|
||||
)
|
||||
await self._process_inbound_event(panel_id, evt, "panel")
|
||||
|
||||
async def _handle_notify_inbox_append(self, payload: Any) -> None:
|
||||
if not isinstance(payload, dict) or payload.get("type") != "message":
|
||||
return
|
||||
detail = payload.get("payload")
|
||||
if not isinstance(detail, dict):
|
||||
return
|
||||
if _str_field(detail, "groupId"):
|
||||
return
|
||||
converse_id = _str_field(detail, "converseId")
|
||||
if not converse_id:
|
||||
return
|
||||
|
||||
session_id = self._session_by_converse.get(converse_id)
|
||||
if not session_id:
|
||||
await self._refresh_sessions_directory(self._ws_ready)
|
||||
session_id = self._session_by_converse.get(converse_id)
|
||||
if not session_id:
|
||||
return
|
||||
|
||||
evt = _make_synthetic_event(
|
||||
message_id=str(detail.get("messageId") or payload.get("_id") or ""),
|
||||
author=str(detail.get("messageAuthor") or ""),
|
||||
content=str(detail.get("messagePlainContent") or detail.get("messageSnippet") or ""),
|
||||
meta={"source": "notify:chat.inbox.append", "converseId": converse_id},
|
||||
group_id="", converse_id=converse_id, timestamp=payload.get("createdAt"),
|
||||
)
|
||||
await self._process_inbound_event(session_id, evt, "session")
|
||||
|
||||
# ---- cursor persistence ------------------------------------------------
|
||||
|
||||
def _mark_session_cursor(self, session_id: str, cursor: int) -> None:
|
||||
if cursor < 0 or cursor < self._session_cursor.get(session_id, 0):
|
||||
return
|
||||
self._session_cursor[session_id] = cursor
|
||||
if not self._cursor_save_task or self._cursor_save_task.done():
|
||||
self._cursor_save_task = asyncio.create_task(self._save_cursor_debounced())
|
||||
|
||||
async def _save_cursor_debounced(self) -> None:
|
||||
await asyncio.sleep(CURSOR_SAVE_DEBOUNCE_S)
|
||||
await self._save_session_cursors()
|
||||
|
||||
async def _load_session_cursors(self) -> None:
|
||||
if not self._cursor_path.exists():
|
||||
return
|
||||
try:
|
||||
data = json.loads(self._cursor_path.read_text("utf-8"))
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to read cursor file: {}", e)
|
||||
return
|
||||
cursors = data.get("cursors") if isinstance(data, dict) else None
|
||||
if isinstance(cursors, dict):
|
||||
for sid, cur in cursors.items():
|
||||
if isinstance(sid, str) and isinstance(cur, int) and cur >= 0:
|
||||
self._session_cursor[sid] = cur
|
||||
|
||||
async def _save_session_cursors(self) -> None:
|
||||
try:
|
||||
self._state_dir.mkdir(parents=True, exist_ok=True)
|
||||
self._cursor_path.write_text(json.dumps({
|
||||
"schemaVersion": 1, "updatedAt": datetime.utcnow().isoformat(),
|
||||
"cursors": self._session_cursor,
|
||||
}, ensure_ascii=False, indent=2) + "\n", "utf-8")
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to save cursor file: {}", e)
|
||||
|
||||
# ---- HTTP helpers ------------------------------------------------------
|
||||
|
||||
async def _post_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
if not self._http:
|
||||
raise RuntimeError("Mochat HTTP client not initialized")
|
||||
url = f"{self.config.base_url.strip().rstrip('/')}{path}"
|
||||
response = await self._http.post(url, headers={
|
||||
"Content-Type": "application/json", "X-Claw-Token": self.config.claw_token,
|
||||
}, json=payload)
|
||||
if not response.is_success:
|
||||
raise RuntimeError(f"Mochat HTTP {response.status_code}: {response.text[:200]}")
|
||||
try:
|
||||
parsed = response.json()
|
||||
except Exception:
|
||||
parsed = response.text
|
||||
if isinstance(parsed, dict) and isinstance(parsed.get("code"), int):
|
||||
if parsed["code"] != 200:
|
||||
msg = str(parsed.get("message") or parsed.get("name") or "request failed")
|
||||
raise RuntimeError(f"Mochat API error: {msg} (code={parsed['code']})")
|
||||
data = parsed.get("data")
|
||||
return data if isinstance(data, dict) else {}
|
||||
return parsed if isinstance(parsed, dict) else {}
|
||||
|
||||
async def _api_send(self, path: str, id_key: str, id_val: str,
|
||||
content: str, reply_to: str | None, group_id: str | None = None) -> dict[str, Any]:
|
||||
"""Unified send helper for session and panel messages."""
|
||||
body: dict[str, Any] = {id_key: id_val, "content": content}
|
||||
if reply_to:
|
||||
body["replyTo"] = reply_to
|
||||
if group_id:
|
||||
body["groupId"] = group_id
|
||||
return await self._post_json(path, body)
|
||||
|
||||
@staticmethod
|
||||
def _read_group_id(metadata: dict[str, Any]) -> str | None:
|
||||
if not isinstance(metadata, dict):
|
||||
return None
|
||||
value = metadata.get("group_id") or metadata.get("groupId")
|
||||
return value.strip() if isinstance(value, str) and value.strip() else None
|
||||
@@ -0,0 +1,822 @@
|
||||
"""Microsoft Teams channel MVP using a tiny built-in HTTP webhook server.
|
||||
|
||||
Scope:
|
||||
- DM-focused MVP
|
||||
- text inbound/outbound
|
||||
- conversation reference persistence
|
||||
- sender allowlist support
|
||||
- optional inbound Bot Framework bearer-token validation
|
||||
- no attachments/cards/polls yet
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager, suppress
|
||||
from dataclasses import dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
try: # pragma: no cover - Windows fallback path
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover
|
||||
fcntl = None
|
||||
|
||||
import httpx
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.config.paths import get_workspace_path
|
||||
from pydantic import BaseModel
|
||||
|
||||
MSTEAMS_AVAILABLE = (
|
||||
importlib.util.find_spec("jwt") is not None
|
||||
and importlib.util.find_spec("cryptography") is not None
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import jwt
|
||||
|
||||
if MSTEAMS_AVAILABLE:
|
||||
import jwt
|
||||
|
||||
MSTEAMS_REF_TTL_DAYS = 30
|
||||
MSTEAMS_WEBCHAT_HOST = "webchat.botframework.com"
|
||||
MSTEAMS_DEFAULT_TRUSTED_SERVICE_URL_HOSTS = [
|
||||
"smba.trafficmanager.net",
|
||||
"smba.infra.gcc.teams.microsoft.com",
|
||||
"smba.infra.gov.teams.microsoft.us",
|
||||
"smba.infra.dod.teams.microsoft.us",
|
||||
"*.botframework.com",
|
||||
]
|
||||
MSTEAMS_REF_META_FILENAME = "msteams_conversations_meta.json"
|
||||
MSTEAMS_REF_LOCK_FILENAME = "msteams_conversations.lock"
|
||||
MSTEAMS_REF_TOUCH_INTERVAL_S = 300
|
||||
|
||||
|
||||
class MSTeamsConfig(BaseModel):
|
||||
"""Microsoft Teams channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
app_id: str = ""
|
||||
app_password: str = ""
|
||||
tenant_id: str = ""
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 3978
|
||||
path: str = "/api/messages"
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
reply_in_thread: bool = True
|
||||
mention_only_response: str = "Hi — what can I help with?"
|
||||
validate_inbound_auth: bool = True
|
||||
ref_ttl_days: int = Field(default=MSTEAMS_REF_TTL_DAYS, ge=1)
|
||||
prune_web_chat_refs: bool = True
|
||||
prune_non_personal_refs: bool = True
|
||||
ref_touch_interval_s: int = Field(default=MSTEAMS_REF_TOUCH_INTERVAL_S, ge=0)
|
||||
trusted_service_url_hosts: list[str] = Field(
|
||||
default_factory=lambda: MSTEAMS_DEFAULT_TRUSTED_SERVICE_URL_HOSTS.copy()
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationRef:
|
||||
"""Minimal stored conversation reference for replies."""
|
||||
|
||||
service_url: str
|
||||
conversation_id: str
|
||||
bot_id: str | None = None
|
||||
activity_id: str | None = None
|
||||
conversation_type: str | None = None
|
||||
tenant_id: str | None = None
|
||||
updated_at: float | None = None
|
||||
|
||||
|
||||
class MSTeamsChannel(BaseChannel):
|
||||
"""Microsoft Teams channel (DM-first MVP)."""
|
||||
|
||||
name = "msteams"
|
||||
display_name = "Microsoft Teams"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return MSTeamsConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = MSTeamsConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: MSTeamsConfig = config
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._server: ThreadingHTTPServer | None = None
|
||||
self._server_thread: threading.Thread | None = None
|
||||
self._http: httpx.AsyncClient | None = None
|
||||
self._token: str | None = None
|
||||
self._token_expires_at: float = 0.0
|
||||
self._botframework_openid_config_url = (
|
||||
"https://login.botframework.com/v1/.well-known/openidconfiguration"
|
||||
)
|
||||
self._botframework_openid_config: dict[str, Any] | None = None
|
||||
self._botframework_openid_config_expires_at: float = 0.0
|
||||
self._botframework_jwks: dict[str, Any] | None = None
|
||||
self._botframework_jwks_expires_at: float = 0.0
|
||||
self._refs_path = get_workspace_path() / "state" / "msteams_conversations.json"
|
||||
self._refs_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._refs_meta_path = self._refs_path.parent / MSTEAMS_REF_META_FILENAME
|
||||
self._refs_lock_path = self._refs_path.parent / MSTEAMS_REF_LOCK_FILENAME
|
||||
self._refs_guard = threading.RLock()
|
||||
self._conversation_refs: dict[str, ConversationRef] = self._load_refs()
|
||||
with self._refs_guard:
|
||||
if self._prune_conversation_refs():
|
||||
self._save_refs_locked(prune=True)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the Teams webhook listener."""
|
||||
if not MSTEAMS_AVAILABLE:
|
||||
self.logger.error("PyJWT not installed. Run: pip install vibe-trading-ai[msteams]")
|
||||
return
|
||||
|
||||
if not self.config.app_id or not self.config.app_password:
|
||||
self.logger.error("app_id/app_password not configured")
|
||||
return
|
||||
|
||||
if not self.config.validate_inbound_auth:
|
||||
self.logger.warning(
|
||||
"Inbound auth validation was explicitly DISABLED in config. "
|
||||
"Anyone who knows the webhook URL can send messages as any user. "
|
||||
"Only disable this for local development or controlled testing."
|
||||
)
|
||||
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._http = httpx.AsyncClient(timeout=30.0)
|
||||
self._running = True
|
||||
|
||||
channel = self
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
if self.path != channel.config.path:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
raw = self.rfile.read(length) if length > 0 else b"{}"
|
||||
payload = json.loads(raw.decode("utf-8"))
|
||||
except Exception as e:
|
||||
channel.logger.warning("Invalid request body: {}", e)
|
||||
self.send_response(400)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
auth_header = self.headers.get("Authorization", "")
|
||||
if channel.config.validate_inbound_auth:
|
||||
try:
|
||||
fut = asyncio.run_coroutine_threadsafe(
|
||||
channel._validate_inbound_auth(auth_header, payload),
|
||||
channel._loop,
|
||||
)
|
||||
fut.result(timeout=15)
|
||||
except Exception as e:
|
||||
channel.logger.warning("Inbound auth validation failed: {}", e)
|
||||
self.send_response(401)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(b'{"error":"unauthorized"}')
|
||||
return
|
||||
try:
|
||||
fut = asyncio.run_coroutine_threadsafe(
|
||||
channel._handle_activity(payload),
|
||||
channel._loop,
|
||||
)
|
||||
fut.result(timeout=15)
|
||||
except Exception as e:
|
||||
channel.logger.warning("Activity handling failed: {}", e)
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"{}")
|
||||
|
||||
def log_message(self, format: str, *args: Any) -> None:
|
||||
return
|
||||
|
||||
self._server = ThreadingHTTPServer((self.config.host, self.config.port), Handler)
|
||||
self._server_thread = threading.Thread(
|
||||
target=self._server.serve_forever,
|
||||
name="vibe-trading-msteams",
|
||||
daemon=True,
|
||||
)
|
||||
self._server_thread.start()
|
||||
|
||||
self.logger.info(
|
||||
"Webhook listening on http://{}:{}{}",
|
||||
self.config.host,
|
||||
self.config.port,
|
||||
self.config.path,
|
||||
)
|
||||
|
||||
while self._running:
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the channel."""
|
||||
self._running = False
|
||||
if self._server:
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
self._server = None
|
||||
if self._server_thread and self._server_thread.is_alive():
|
||||
self._server_thread.join(timeout=2)
|
||||
self._server_thread = None
|
||||
if self._http:
|
||||
await self._http.aclose()
|
||||
self._http = None
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a plain text reply into an existing Teams conversation."""
|
||||
if not self._http:
|
||||
raise RuntimeError("MSTeams HTTP client not initialized")
|
||||
|
||||
ref = self._conversation_refs.get(str(msg.chat_id))
|
||||
if not ref:
|
||||
raise RuntimeError(f"MSTeams conversation ref not found for chat_id={msg.chat_id}")
|
||||
|
||||
if not self._is_trusted_service_url(ref.service_url):
|
||||
raise RuntimeError(
|
||||
f"MSTeams conversation ref has untrusted service_url for chat_id={msg.chat_id}"
|
||||
)
|
||||
|
||||
token = await self._get_access_token()
|
||||
base_url = f"{ref.service_url.rstrip('/')}/v3/conversations/{ref.conversation_id}/activities"
|
||||
use_thread_reply = self.config.reply_in_thread and bool(ref.activity_id)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"type": "message",
|
||||
"text": msg.content or " ",
|
||||
}
|
||||
if use_thread_reply:
|
||||
payload["replyToId"] = ref.activity_id
|
||||
|
||||
try:
|
||||
resp = await self._http.post(base_url, headers=headers, json=payload)
|
||||
resp.raise_for_status()
|
||||
self.logger.info("Message sent to {}", ref.conversation_id)
|
||||
self._touch_conversation_ref(str(msg.chat_id), persist=True)
|
||||
except Exception:
|
||||
self.logger.exception("Send failed")
|
||||
raise
|
||||
|
||||
async def _handle_activity(self, activity: dict[str, Any]) -> None:
|
||||
"""Handle inbound Teams/Bot Framework activity."""
|
||||
if activity.get("type") != "message":
|
||||
return
|
||||
|
||||
conversation = activity.get("conversation") or {}
|
||||
from_user = activity.get("from") or {}
|
||||
recipient = activity.get("recipient") or {}
|
||||
channel_data = activity.get("channelData") or {}
|
||||
|
||||
sender_id = str(from_user.get("aadObjectId") or from_user.get("id") or "").strip()
|
||||
conversation_id = str(conversation.get("id") or "").strip()
|
||||
service_url = str(activity.get("serviceUrl") or "").strip()
|
||||
activity_id = str(activity.get("id") or "").strip()
|
||||
conversation_type = str(conversation.get("conversationType") or "").strip()
|
||||
|
||||
if not sender_id or not conversation_id or not service_url:
|
||||
return
|
||||
|
||||
if not self._is_trusted_service_url(service_url):
|
||||
self.logger.warning(
|
||||
"Ignoring MSTeams activity with untrusted serviceUrl host: {}",
|
||||
service_url,
|
||||
)
|
||||
return
|
||||
|
||||
if recipient.get("id") and from_user.get("id") == recipient.get("id"):
|
||||
return
|
||||
|
||||
# DM-only MVP: ignore group/channel traffic for now
|
||||
if conversation_type and conversation_type not in ("personal", ""):
|
||||
self.logger.debug("Ignoring non-DM conversation {}", conversation_type)
|
||||
return
|
||||
|
||||
text = self._sanitize_inbound_text(activity)
|
||||
if not text:
|
||||
text = self.config.mention_only_response.strip()
|
||||
if not text:
|
||||
self.logger.debug("Ignoring empty message after Teams text sanitization")
|
||||
return
|
||||
|
||||
if not self.is_allowed(sender_id):
|
||||
self.logger.warning(
|
||||
"Access denied for sender {} on channel {}. "
|
||||
"Add them to allowFrom list in config to grant access.",
|
||||
sender_id, self.name,
|
||||
)
|
||||
return
|
||||
|
||||
with self._refs_guard:
|
||||
self._conversation_refs[conversation_id] = ConversationRef(
|
||||
service_url=service_url,
|
||||
conversation_id=conversation_id,
|
||||
bot_id=str(recipient.get("id") or "") or None,
|
||||
activity_id=activity_id or None,
|
||||
conversation_type=conversation_type or None,
|
||||
tenant_id=str((channel_data.get("tenant") or {}).get("id") or "") or None,
|
||||
updated_at=time.time(),
|
||||
)
|
||||
self._save_refs_locked()
|
||||
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=conversation_id,
|
||||
content=text,
|
||||
metadata={
|
||||
"msteams": {
|
||||
"activity_id": activity_id,
|
||||
"conversation_id": conversation_id,
|
||||
"conversation_type": conversation_type or "personal",
|
||||
"from_name": from_user.get("name"),
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
def _sanitize_inbound_text(self, activity: dict[str, Any]) -> str:
|
||||
"""Extract the user-authored text from a Teams activity."""
|
||||
text = str(activity.get("text") or "")
|
||||
text = self._strip_possible_bot_mention(text)
|
||||
text = self._normalize_html_whitespace(text)
|
||||
|
||||
channel_data = activity.get("channelData") or {}
|
||||
reply_to_id = str(activity.get("replyToId") or "").strip()
|
||||
normalized_preview = html.unescape(text).replace("&rsquo", "’").strip()
|
||||
normalized_preview = normalized_preview.replace("\xa0", " ")
|
||||
normalized_preview = normalized_preview.replace("\r\n", "\n").replace("\r", "\n")
|
||||
preview_lines = [line.strip() for line in normalized_preview.split("\n")]
|
||||
while preview_lines and not preview_lines[0]:
|
||||
preview_lines.pop(0)
|
||||
first_line = preview_lines[0] if preview_lines else ""
|
||||
looks_like_quote_wrapper = first_line.lower().startswith("replying to ") or first_line.startswith("Reply wrapper")
|
||||
|
||||
if reply_to_id or channel_data.get("messageType") == "reply" or looks_like_quote_wrapper:
|
||||
text = self._normalize_teams_reply_quote(text)
|
||||
|
||||
return text.strip()
|
||||
|
||||
def _strip_possible_bot_mention(self, text: str) -> str:
|
||||
"""Remove simple Teams mention markup from message text."""
|
||||
cleaned = re.sub(r"<at\b[^>]*>.*?</at>", " ", text, flags=re.IGNORECASE | re.DOTALL)
|
||||
cleaned = re.sub(r"[^\S\r\n]+", " ", cleaned)
|
||||
cleaned = re.sub(r"(?:\r?\n){3,}", "\n\n", cleaned)
|
||||
return cleaned.strip()
|
||||
|
||||
def _normalize_html_whitespace(self, text: str) -> str:
|
||||
"""Normalize common HTML whitespace/entities from Teams into plain text spacing."""
|
||||
normalized = html.unescape(text).replace("&rsquo", "’")
|
||||
normalized = normalized.replace("\xa0", " ")
|
||||
return normalized
|
||||
|
||||
def _normalize_teams_reply_quote(self, text: str) -> str:
|
||||
"""Normalize Teams quoted replies into a compact structured form."""
|
||||
cleaned = self._normalize_html_whitespace(text).strip()
|
||||
if not cleaned:
|
||||
return ""
|
||||
|
||||
normalized_newlines = cleaned.replace("\r\n", "\n").replace("\r", "\n")
|
||||
lines = [line.strip() for line in normalized_newlines.split("\n")]
|
||||
while lines and not lines[0]:
|
||||
lines.pop(0)
|
||||
|
||||
# Observed native Teams reply wrapper:
|
||||
# Replying to Bob Smith
|
||||
# actual reply text
|
||||
if len(lines) >= 2 and lines[0].lower().startswith("replying to "):
|
||||
quoted = lines[0][len("replying to ") :].strip(" :")
|
||||
reply = "\n".join(lines[1:]).strip()
|
||||
return self._format_reply_with_quote(quoted, reply)
|
||||
|
||||
# Observed reply wrapper where the quoted content is surfaced after a
|
||||
# synthetic "Reply wrapper" header, sometimes with a blank line separating quote
|
||||
# and reply, and sometimes as a compact line-based fallback shape.
|
||||
if lines and lines[0].strip().startswith("Reply wrapper"):
|
||||
body = normalized_newlines.split("\n", 1)[1] if "\n" in normalized_newlines else ""
|
||||
body = body.lstrip()
|
||||
parts = re.split(r"\n\s*\n", body, maxsplit=1)
|
||||
if len(parts) == 2:
|
||||
quoted = re.sub(r"\s+", " ", parts[0]).strip()
|
||||
reply = re.sub(r"\s+", " ", parts[1]).strip()
|
||||
if quoted or reply:
|
||||
return self._format_reply_with_quote(quoted, reply)
|
||||
|
||||
body_lines = [line.strip() for line in body.split("\n") if line.strip()]
|
||||
if body_lines:
|
||||
quoted = " ".join(body_lines[:-1]).strip()
|
||||
reply = body_lines[-1].strip()
|
||||
if quoted and reply:
|
||||
return self._format_reply_with_quote(quoted, reply)
|
||||
|
||||
# Observed compact fallback where the relay flattens quote and reply into
|
||||
# a single line after the synthetic Reply wrapper prefix.
|
||||
compact = re.sub(r"\s+", " ", normalized_newlines).strip()
|
||||
if compact.startswith("Reply wrapper "):
|
||||
compact = compact[len("Reply wrapper ") :].strip()
|
||||
for boundary in (". ", "! ", "? ", "… "):
|
||||
idx = compact.rfind(boundary)
|
||||
if idx == -1:
|
||||
continue
|
||||
quoted = compact[: idx + 1].strip()
|
||||
reply = compact[idx + len(boundary) :].strip()
|
||||
if quoted and reply and len(reply) <= 160:
|
||||
return self._format_reply_with_quote(quoted, reply)
|
||||
|
||||
return cleaned
|
||||
|
||||
def _format_reply_with_quote(self, quoted: str, reply: str) -> str:
|
||||
"""Format a reply-with-context message for the model without Teams wrapper noise."""
|
||||
quoted = quoted.strip()
|
||||
reply = reply.strip()
|
||||
if quoted and reply:
|
||||
return f"User is replying to: {quoted}\nUser reply: {reply}"
|
||||
if reply:
|
||||
return reply
|
||||
return quoted
|
||||
|
||||
async def _validate_inbound_auth(self, auth_header: str, activity: dict[str, Any]) -> None:
|
||||
"""Validate inbound Bot Framework bearer token."""
|
||||
if not MSTEAMS_AVAILABLE:
|
||||
raise RuntimeError("PyJWT not installed. Run: pip install vibe-trading-ai[msteams]")
|
||||
|
||||
if not auth_header.lower().startswith("bearer "):
|
||||
raise ValueError("missing bearer token")
|
||||
|
||||
token = auth_header.split(" ", 1)[1].strip()
|
||||
if not token:
|
||||
raise ValueError("empty bearer token")
|
||||
|
||||
header = jwt.get_unverified_header(token)
|
||||
kid = str(header.get("kid") or "").strip()
|
||||
if not kid:
|
||||
raise ValueError("missing token kid")
|
||||
|
||||
jwks = await self._get_botframework_jwks()
|
||||
keys = jwks.get("keys") or []
|
||||
jwk = next((key for key in keys if key.get("kid") == kid), None)
|
||||
if not jwk:
|
||||
raise ValueError(f"signing key not found for kid={kid}")
|
||||
|
||||
public_key = jwt.algorithms.RSAAlgorithm.from_jwk(json.dumps(jwk))
|
||||
claims = jwt.decode(
|
||||
token,
|
||||
key=public_key,
|
||||
algorithms=["RS256"],
|
||||
audience=self.config.app_id,
|
||||
issuer="https://api.botframework.com",
|
||||
options={
|
||||
"require": ["exp", "nbf", "iss", "aud"],
|
||||
},
|
||||
)
|
||||
|
||||
claim_service_url = str(
|
||||
claims.get("serviceurl") or claims.get("serviceUrl") or "",
|
||||
).strip()
|
||||
activity_service_url = str(activity.get("serviceUrl") or "").strip()
|
||||
if claim_service_url and activity_service_url and claim_service_url != activity_service_url:
|
||||
raise ValueError("serviceUrl claim mismatch")
|
||||
|
||||
async def _get_botframework_openid_config(self) -> dict[str, Any]:
|
||||
"""Fetch and cache Bot Framework OpenID configuration."""
|
||||
|
||||
now = time.time()
|
||||
if self._botframework_openid_config and now < self._botframework_openid_config_expires_at:
|
||||
return self._botframework_openid_config
|
||||
|
||||
if not self._http:
|
||||
raise RuntimeError("MSTeams HTTP client not initialized")
|
||||
|
||||
resp = await self._http.get(self._botframework_openid_config_url)
|
||||
resp.raise_for_status()
|
||||
self._botframework_openid_config = resp.json()
|
||||
self._botframework_openid_config_expires_at = now + 3600
|
||||
return self._botframework_openid_config
|
||||
|
||||
async def _get_botframework_jwks(self) -> dict[str, Any]:
|
||||
"""Fetch and cache Bot Framework JWKS."""
|
||||
|
||||
now = time.time()
|
||||
if self._botframework_jwks and now < self._botframework_jwks_expires_at:
|
||||
return self._botframework_jwks
|
||||
|
||||
if not self._http:
|
||||
raise RuntimeError("MSTeams HTTP client not initialized")
|
||||
|
||||
openid_config = await self._get_botframework_openid_config()
|
||||
jwks_uri = str(openid_config.get("jwks_uri") or "").strip()
|
||||
if not jwks_uri:
|
||||
raise RuntimeError("Bot Framework OpenID config missing jwks_uri")
|
||||
|
||||
resp = await self._http.get(jwks_uri)
|
||||
resp.raise_for_status()
|
||||
self._botframework_jwks = resp.json()
|
||||
self._botframework_jwks_expires_at = now + 3600
|
||||
return self._botframework_jwks
|
||||
|
||||
@staticmethod
|
||||
def _safe_float(value: Any) -> float | None:
|
||||
try:
|
||||
out = float(value)
|
||||
if out > 0:
|
||||
return out
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return None
|
||||
|
||||
def _normalize_ref_record(self, value: Any) -> ConversationRef | None:
|
||||
"""Normalize a stored ref record from legacy/current schema."""
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
service_url = str(value.get("service_url") or "").strip()
|
||||
conversation_id = str(value.get("conversation_id") or "").strip()
|
||||
if not service_url or not conversation_id:
|
||||
return None
|
||||
return ConversationRef(
|
||||
service_url=service_url,
|
||||
conversation_id=conversation_id,
|
||||
bot_id=str(value.get("bot_id") or "") or None,
|
||||
activity_id=str(value.get("activity_id") or "") or None,
|
||||
conversation_type=str(value.get("conversation_type") or "") or None,
|
||||
tenant_id=str(value.get("tenant_id") or "") or None,
|
||||
updated_at=self._safe_float(value.get("updated_at")),
|
||||
)
|
||||
|
||||
def _load_refs_raw(self) -> tuple[dict[str, Any], dict[str, Any], bool]:
|
||||
"""Load raw refs/main+meta JSON payloads."""
|
||||
main_data: dict[str, Any] = {}
|
||||
meta_data: dict[str, Any] = {}
|
||||
meta_exists = self._refs_meta_path.exists()
|
||||
|
||||
if self._refs_path.exists():
|
||||
try:
|
||||
loaded = json.loads(self._refs_path.read_text(encoding="utf-8"))
|
||||
if isinstance(loaded, dict):
|
||||
main_data = loaded
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to load conversation refs: {}", e)
|
||||
|
||||
if meta_exists:
|
||||
try:
|
||||
loaded_meta = json.loads(self._refs_meta_path.read_text(encoding="utf-8"))
|
||||
if isinstance(loaded_meta, dict):
|
||||
meta_data = loaded_meta
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to load conversation refs metadata: {}", e)
|
||||
|
||||
return main_data, meta_data, meta_exists
|
||||
|
||||
def _load_refs_from_disk(self) -> dict[str, ConversationRef]:
|
||||
"""Load refs from disk with compatibility fallback for legacy layouts."""
|
||||
main_data, meta_data, meta_exists = self._load_refs_raw()
|
||||
if not main_data:
|
||||
return {}
|
||||
|
||||
out: dict[str, ConversationRef] = {}
|
||||
now = time.time()
|
||||
for key, value in main_data.items():
|
||||
ref = self._normalize_ref_record(value)
|
||||
if not ref:
|
||||
continue
|
||||
|
||||
meta_entry = meta_data.get(key) if isinstance(meta_data, dict) else None
|
||||
meta_ts = None
|
||||
if isinstance(meta_entry, dict):
|
||||
meta_ts = self._safe_float(meta_entry.get("updated_at"))
|
||||
elif meta_entry is not None:
|
||||
meta_ts = self._safe_float(meta_entry)
|
||||
|
||||
if meta_ts is not None:
|
||||
ref.updated_at = meta_ts
|
||||
elif not meta_exists:
|
||||
# First run after introducing meta sidecar: keep legacy refs alive
|
||||
# by initializing timestamps to "now" instead of purging immediately.
|
||||
ref.updated_at = now
|
||||
elif ref.updated_at is None:
|
||||
ref.updated_at = now
|
||||
|
||||
out[key] = ref
|
||||
return out
|
||||
|
||||
def _load_refs(self) -> dict[str, ConversationRef]:
|
||||
"""Load stored conversation references."""
|
||||
return self._load_refs_from_disk()
|
||||
|
||||
@contextmanager
|
||||
def _refs_file_lock(self):
|
||||
"""Cross-process lock while merging and writing refs state."""
|
||||
self._refs_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
lock_fp = self._refs_lock_path.open("a+", encoding="utf-8")
|
||||
try:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(lock_fp.fileno(), fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(lock_fp.fileno(), fcntl.LOCK_UN)
|
||||
finally:
|
||||
lock_fp.close()
|
||||
|
||||
def _is_webchat_service_url(self, service_url: str) -> bool:
|
||||
"""Return True when service URL points to unsupported Bot Framework Web Chat."""
|
||||
normalized = service_url.strip()
|
||||
if not normalized:
|
||||
return False
|
||||
host = (urlparse(normalized).hostname or "").strip().lower()
|
||||
if host:
|
||||
return host == MSTEAMS_WEBCHAT_HOST or host.endswith(f".{MSTEAMS_WEBCHAT_HOST}")
|
||||
return MSTEAMS_WEBCHAT_HOST in normalized.lower()
|
||||
|
||||
def _is_trusted_service_url(self, service_url: str) -> bool:
|
||||
"""Return True for HTTPS Bot Framework service URLs trusted for bearer replies."""
|
||||
parsed = urlparse(service_url.strip())
|
||||
if parsed.scheme.lower() != "https":
|
||||
return False
|
||||
|
||||
host = (parsed.hostname or "").strip().lower().rstrip(".")
|
||||
if not host:
|
||||
return False
|
||||
|
||||
for pattern in self.config.trusted_service_url_hosts:
|
||||
trusted_host = str(pattern or "").strip().lower().rstrip(".")
|
||||
if not trusted_host:
|
||||
continue
|
||||
if trusted_host.startswith("*."):
|
||||
suffix = trusted_host[1:]
|
||||
if host.endswith(suffix) and host != suffix.lstrip("."):
|
||||
return True
|
||||
continue
|
||||
if host == trusted_host:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _prune_conversation_refs(self, *, now: float | None = None) -> bool:
|
||||
"""Remove stale and unsupported conversation refs from memory."""
|
||||
if not self._conversation_refs:
|
||||
return False
|
||||
|
||||
now_ts = time.time() if now is None else now
|
||||
ttl_days = int(self.config.ref_ttl_days)
|
||||
stale_before = now_ts - (ttl_days * 24 * 60 * 60)
|
||||
keys_to_drop: list[str] = []
|
||||
|
||||
for key, ref in self._conversation_refs.items():
|
||||
if not self._is_trusted_service_url(ref.service_url):
|
||||
keys_to_drop.append(key)
|
||||
continue
|
||||
|
||||
if self.config.prune_web_chat_refs and self._is_webchat_service_url(ref.service_url):
|
||||
keys_to_drop.append(key)
|
||||
continue
|
||||
|
||||
conv_type = str(ref.conversation_type or "").strip().lower()
|
||||
if self.config.prune_non_personal_refs and conv_type and conv_type != "personal":
|
||||
keys_to_drop.append(key)
|
||||
continue
|
||||
|
||||
try:
|
||||
updated_at = float(ref.updated_at) if ref.updated_at is not None else 0.0
|
||||
except (TypeError, ValueError):
|
||||
updated_at = 0.0
|
||||
if updated_at <= 0 or updated_at < stale_before:
|
||||
keys_to_drop.append(key)
|
||||
|
||||
if not keys_to_drop:
|
||||
return False
|
||||
|
||||
for key in keys_to_drop:
|
||||
self._conversation_refs.pop(key, None)
|
||||
self.logger.info(
|
||||
"Pruned {} stale/unsupported conversation refs (ttl={} days)",
|
||||
len(keys_to_drop),
|
||||
ttl_days,
|
||||
)
|
||||
return True
|
||||
|
||||
def _merge_refs_from_disk_locked(self) -> None:
|
||||
"""Merge disk refs into memory to reduce lost updates across processes."""
|
||||
disk_refs = self._load_refs_from_disk()
|
||||
for key, disk_ref in disk_refs.items():
|
||||
mem_ref = self._conversation_refs.get(key)
|
||||
if mem_ref is None:
|
||||
self._conversation_refs[key] = disk_ref
|
||||
continue
|
||||
disk_ts = self._safe_float(disk_ref.updated_at) or 0.0
|
||||
mem_ts = self._safe_float(mem_ref.updated_at) or 0.0
|
||||
if disk_ts > mem_ts:
|
||||
self._conversation_refs[key] = disk_ref
|
||||
|
||||
def _touch_conversation_ref(self, chat_id: str, *, persist: bool = False) -> None:
|
||||
"""Refresh updated_at for an active ref to keep it from expiring while used."""
|
||||
with self._refs_guard:
|
||||
ref = self._conversation_refs.get(str(chat_id))
|
||||
if not ref:
|
||||
return
|
||||
now = time.time()
|
||||
prev = self._safe_float(ref.updated_at) or 0.0
|
||||
min_interval = max(0, int(self.config.ref_touch_interval_s))
|
||||
if min_interval > 0 and prev > 0 and now - prev < min_interval:
|
||||
return
|
||||
ref.updated_at = now
|
||||
if persist:
|
||||
self._save_refs_locked()
|
||||
|
||||
def _write_json_atomically(self, path, data: dict[str, Any]) -> None:
|
||||
"""Write refs JSON atomically to reduce corruption risk during crashes."""
|
||||
payload = json.dumps(data, indent=2)
|
||||
tmp_path: str | None = None
|
||||
try:
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
dir=str(path.parent),
|
||||
prefix=f"{path.name}.",
|
||||
suffix=".tmp",
|
||||
)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(payload)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
finally:
|
||||
if tmp_path and os.path.exists(tmp_path):
|
||||
with suppress(OSError):
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def _save_refs_locked(self, *, prune: bool = True) -> None:
|
||||
"""Persist conversation references (caller must hold _refs_guard)."""
|
||||
try:
|
||||
with self._refs_file_lock():
|
||||
self._merge_refs_from_disk_locked()
|
||||
if prune:
|
||||
self._prune_conversation_refs()
|
||||
refs_data = {
|
||||
key: {
|
||||
"service_url": ref.service_url,
|
||||
"conversation_id": ref.conversation_id,
|
||||
"bot_id": ref.bot_id,
|
||||
"activity_id": ref.activity_id,
|
||||
"conversation_type": ref.conversation_type,
|
||||
"tenant_id": ref.tenant_id,
|
||||
}
|
||||
for key, ref in self._conversation_refs.items()
|
||||
}
|
||||
refs_meta = {
|
||||
key: {
|
||||
"updated_at": self._safe_float(ref.updated_at),
|
||||
}
|
||||
for key, ref in self._conversation_refs.items()
|
||||
}
|
||||
self._write_json_atomically(self._refs_path, refs_data)
|
||||
self._write_json_atomically(self._refs_meta_path, refs_meta)
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to save conversation refs: {}", e)
|
||||
|
||||
def _save_refs(self, *, prune: bool = True) -> None:
|
||||
"""Persist conversation references."""
|
||||
with self._refs_guard:
|
||||
self._save_refs_locked(prune=prune)
|
||||
|
||||
async def _get_access_token(self) -> str:
|
||||
"""Fetch an access token for Bot Framework / Azure Bot auth."""
|
||||
|
||||
now = time.time()
|
||||
if self._token and now < self._token_expires_at - 60:
|
||||
return self._token
|
||||
|
||||
if not self._http:
|
||||
raise RuntimeError("MSTeams HTTP client not initialized")
|
||||
|
||||
tenant = (self.config.tenant_id or "").strip() or "botframework.com"
|
||||
token_url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
|
||||
data = {
|
||||
"grant_type": "client_credentials",
|
||||
"client_id": self.config.app_id,
|
||||
"client_secret": self.config.app_password,
|
||||
"scope": "https://api.botframework.com/.default",
|
||||
}
|
||||
resp = await self._http.post(token_url, data=data)
|
||||
resp.raise_for_status()
|
||||
payload = resp.json()
|
||||
self._token = payload["access_token"]
|
||||
self._token_expires_at = now + int(payload.get("expires_in", 3600))
|
||||
return self._token
|
||||
@@ -0,0 +1,579 @@
|
||||
"""Napcat (OneBot v11) channel for QQ, over WebSocket."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
import aiohttp
|
||||
import logging; logger = logging.getLogger(__name__)
|
||||
from pydantic import Field
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
from websockets.asyncio.client import connect as ws_connect
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_media_dir
|
||||
from pydantic import BaseModel
|
||||
from src.channels.utils import validate_url_target
|
||||
from src.channels.utils import safe_filename
|
||||
|
||||
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
||||
_ACTION_TIMEOUT = 20.0
|
||||
|
||||
|
||||
# `"mention"` (only @mentions / replies) | `"open"` (every message) | float p
|
||||
# in [0, 1]: mentions/replies always reply; other messages reply with probability
|
||||
# p. 0.0 ≡ "mention", 1.0 ≡ "open".
|
||||
GroupPolicy = Literal["mention", "open"] | Annotated[float, Field(ge=0.0, le=1.0)]
|
||||
|
||||
|
||||
class NapcatConfig(BaseModel):
|
||||
"""Napcat (OneBot v11) channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
ws_url: str = "ws://127.0.0.1:3001"
|
||||
access_token: str = ""
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
group_policy: GroupPolicy = "mention"
|
||||
# Per-group overrides keyed by stringified group_id, e.g. {"123456": "open"}.
|
||||
# Falls back to `group_policy` when a group_id isn't listed.
|
||||
group_policy_overrides: dict[str, GroupPolicy] = Field(default_factory=dict)
|
||||
welcome_new_members: bool = True
|
||||
# Hard cap for inbound image downloads. Bigger images are dropped.
|
||||
max_image_bytes: int = Field(default=20 * 1024 * 1024, ge=1)
|
||||
|
||||
|
||||
class NapcatChannel(BaseChannel):
|
||||
"""Napcat / OneBot v11 channel."""
|
||||
|
||||
name = "napcat"
|
||||
display_name = "Napcat (QQ)"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return NapcatConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = NapcatConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: NapcatConfig = config
|
||||
|
||||
self._ws: ClientConnection | None = None
|
||||
self._http: aiohttp.ClientSession | None = None
|
||||
self._media_root: Path = get_media_dir("napcat")
|
||||
self._self_id: int | None = None
|
||||
self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
|
||||
self._processed_ids: deque[int] = deque(maxlen=2000)
|
||||
self._bot_outbound_ids: deque[int] = deque(maxlen=2000)
|
||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def start(self) -> None:
|
||||
if not self.config.ws_url:
|
||||
logger.error("napcat: ws_url not configured")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._http = aiohttp.ClientSession(timeout=_DOWNLOAD_TIMEOUT)
|
||||
|
||||
backoff = iter((5, 10)) # then 30s forever
|
||||
while self._running:
|
||||
try:
|
||||
await self._run_once()
|
||||
backoff = iter((5, 10)) # reset after a clean session
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning("napcat: connection lost: {}", e)
|
||||
if self._running:
|
||||
await asyncio.sleep(next(backoff, 30))
|
||||
|
||||
async def _run_once(self) -> None:
|
||||
headers = []
|
||||
if self.config.access_token:
|
||||
headers.append(("Authorization", f"Bearer {self.config.access_token}"))
|
||||
|
||||
logger.info("napcat: connecting to {}", self.config.ws_url)
|
||||
async with ws_connect(self.config.ws_url, additional_headers=headers) as ws:
|
||||
self._ws = ws
|
||||
logger.info("napcat: connected")
|
||||
try:
|
||||
# Validate the connection before entering the dispatch loop.
|
||||
# Napcat may interleave meta_event frames before our echo
|
||||
# response, so dispatch any non-matching frames as we go.
|
||||
echo = uuid.uuid4().hex
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{"action": "get_login_info", "params": {}, "echo": echo},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
)
|
||||
deadline = asyncio.get_running_loop().time() + _ACTION_TIMEOUT
|
||||
while True:
|
||||
remaining = deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
raise asyncio.TimeoutError("get_login_info timed out")
|
||||
raw = await asyncio.wait_for(ws.recv(), timeout=remaining)
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(payload, dict) and payload.get("echo") == echo:
|
||||
data = payload.get("data") or {}
|
||||
logger.info(
|
||||
"napcat: logged in as {} (user_id={})",
|
||||
data.get("nickname"),
|
||||
data.get("user_id"),
|
||||
)
|
||||
break
|
||||
await self._dispatch_frame(raw)
|
||||
|
||||
async for raw in ws:
|
||||
await self._dispatch_frame(raw)
|
||||
finally:
|
||||
self._ws = None
|
||||
self._fail_pending(RuntimeError("napcat: websocket disconnected"))
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._running = False
|
||||
if self._ws is not None:
|
||||
try:
|
||||
await self._ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._ws = None
|
||||
if self._http is not None:
|
||||
try:
|
||||
await self._http.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._http = None
|
||||
self._fail_pending(RuntimeError("napcat: stopped"))
|
||||
tasks = list(self._background_tasks)
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._background_tasks.clear()
|
||||
|
||||
def _fail_pending(self, err: BaseException) -> None:
|
||||
for fut in self._pending.values():
|
||||
if not fut.done():
|
||||
fut.set_exception(err)
|
||||
self._pending.clear()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Frame dispatch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _dispatch_frame(self, raw: str | bytes) -> None:
|
||||
# logger.debug("dispatch frame {}", raw)
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.debug("napcat: dropping non-JSON frame")
|
||||
return
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
|
||||
# Action response: identified by `echo` and absence of post_type.
|
||||
if "echo" in payload and payload.get("post_type") is None:
|
||||
echo = payload.get("echo")
|
||||
fut = self._pending.pop(echo, None) if isinstance(echo, str) else None
|
||||
if fut and not fut.done():
|
||||
fut.set_result(payload)
|
||||
return
|
||||
|
||||
if (sid := payload.get("self_id")) is not None:
|
||||
try:
|
||||
self._self_id = int(sid)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
post_type = payload.get("post_type")
|
||||
if post_type == "message":
|
||||
self._create_background_task(self._on_message(payload), "message")
|
||||
elif post_type == "notice":
|
||||
self._create_background_task(self._on_notice(payload), "notice")
|
||||
|
||||
def _create_background_task(self, coro: Any, kind: str) -> None:
|
||||
task = asyncio.create_task(coro)
|
||||
self._background_tasks.add(task)
|
||||
|
||||
def _done(done: asyncio.Task[None]) -> None:
|
||||
self._background_tasks.discard(done)
|
||||
try:
|
||||
done.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning("napcat: {} handler failed: {}", kind, e)
|
||||
|
||||
task.add_done_callback(_done)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Inbound: messages
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _on_message(self, ev: dict[str, Any]) -> None:
|
||||
msg_id = ev.get("message_id")
|
||||
if isinstance(msg_id, int):
|
||||
if msg_id in self._processed_ids:
|
||||
return
|
||||
self._processed_ids.append(msg_id)
|
||||
|
||||
message_type = ev.get("message_type")
|
||||
user_id = ev.get("user_id")
|
||||
if user_id is None or message_type not in ("group", "private"):
|
||||
return
|
||||
|
||||
segments = self._normalize_segments(ev.get("message"))
|
||||
text, images, mentioned_self, reply_to_id = self._parse_segments(segments)
|
||||
|
||||
media_paths: list[str] = []
|
||||
for info in images:
|
||||
if local := await self._download_image(info):
|
||||
media_paths.append(local)
|
||||
|
||||
sender = ev.get("sender") or {}
|
||||
nickname = sender.get("card") or sender.get("nickname")
|
||||
|
||||
if message_type == "group":
|
||||
group_id = ev.get("group_id")
|
||||
if group_id is None:
|
||||
return
|
||||
|
||||
replying_to_bot = (
|
||||
isinstance(reply_to_id, int) and reply_to_id in self._bot_outbound_ids
|
||||
)
|
||||
if not self._should_reply_in_group(
|
||||
group_id=group_id,
|
||||
mentioned_self=mentioned_self,
|
||||
replying_to_bot=replying_to_bot,
|
||||
):
|
||||
return
|
||||
|
||||
chat_id = f"group:{group_id}"
|
||||
content = self._format_group_content(
|
||||
text=text,
|
||||
nickname=nickname,
|
||||
user_id=user_id,
|
||||
)
|
||||
else:
|
||||
chat_id = f"private:{user_id}"
|
||||
content = text
|
||||
|
||||
if not content and not media_paths:
|
||||
return
|
||||
|
||||
await self._handle_message(
|
||||
sender_id=str(user_id),
|
||||
chat_id=chat_id,
|
||||
content=content,
|
||||
media=media_paths or None,
|
||||
metadata={
|
||||
"message_id": msg_id,
|
||||
"is_group": message_type == "group",
|
||||
"nickname": nickname,
|
||||
"reply_to": reply_to_id,
|
||||
},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_segments(message: Any) -> list[dict[str, Any]]:
|
||||
# Napcat defaults to array format. Treat raw strings as a single text
|
||||
# segment rather than parsing CQ codes — that path is fragile and
|
||||
# users can configure napcat to emit arrays.
|
||||
if isinstance(message, list):
|
||||
return [seg for seg in message if isinstance(seg, dict)]
|
||||
if isinstance(message, str) and message:
|
||||
return [{"type": "text", "data": {"text": message}}]
|
||||
return []
|
||||
|
||||
def _parse_segments(
|
||||
self, segments: list[dict[str, Any]]
|
||||
) -> tuple[str, list[dict[str, Any]], bool, int | None]:
|
||||
parts: list[str] = []
|
||||
images: list[dict[str, Any]] = []
|
||||
mentioned_self = False
|
||||
reply_to: int | None = None
|
||||
self_id_str = str(self._self_id) if self._self_id is not None else None
|
||||
|
||||
for seg in segments:
|
||||
stype = seg.get("type")
|
||||
data = seg.get("data") or {}
|
||||
if stype == "text":
|
||||
if txt := data.get("text"):
|
||||
parts.append(str(txt))
|
||||
elif stype == "image":
|
||||
# OneBot exposes the downloadable image at `url`. Napcat
|
||||
# additionally provides `file` (e.g. <md5>.png) and
|
||||
# `file_size` (bytes, sometimes a string).
|
||||
url = data.get("url")
|
||||
if isinstance(url, str) and url.startswith(("http://", "https://")):
|
||||
images.append(
|
||||
{
|
||||
"url": url,
|
||||
"file": data.get("file"),
|
||||
"file_size": data.get("file_size"),
|
||||
}
|
||||
)
|
||||
else:
|
||||
logger.warning("napcat: received invalid image url: {}", url)
|
||||
elif stype == "at":
|
||||
qq = str(data.get("qq", ""))
|
||||
if self_id_str and qq == self_id_str:
|
||||
mentioned_self = True
|
||||
else:
|
||||
parts.append(f"@{qq}")
|
||||
elif stype == "reply":
|
||||
rid = data.get("id")
|
||||
try:
|
||||
reply_to = int(rid) if rid is not None else None
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
elif stype == "face":
|
||||
parts.append(f"[face:{data.get('id', '')}]")
|
||||
|
||||
text = " ".join(p.strip() for p in parts if p.strip()).strip()
|
||||
return text, images, mentioned_self, reply_to
|
||||
|
||||
def _should_reply_in_group(
|
||||
self, *, group_id: Any, mentioned_self: bool, replying_to_bot: bool
|
||||
) -> bool:
|
||||
if mentioned_self or replying_to_bot:
|
||||
return True
|
||||
policy = self.config.group_policy_overrides.get(str(group_id), self.config.group_policy)
|
||||
if policy == "open":
|
||||
return True
|
||||
if policy == "mention":
|
||||
return False
|
||||
# Probability case: float in [0.0, 1.0].
|
||||
return random.random() < float(policy)
|
||||
|
||||
@staticmethod
|
||||
def _format_group_content(
|
||||
*,
|
||||
text: str,
|
||||
nickname: str,
|
||||
user_id: Any,
|
||||
) -> str:
|
||||
label = nickname or str(user_id)
|
||||
return f"{label}: {text}"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Inbound: notices (member joined etc.)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _on_notice(self, ev: dict[str, Any]) -> None:
|
||||
if ev.get("notice_type") != "group_increase" or not self.config.welcome_new_members:
|
||||
return
|
||||
|
||||
group_id = ev.get("group_id")
|
||||
user_id = ev.get("user_id")
|
||||
if group_id is None or user_id is None:
|
||||
return
|
||||
|
||||
try:
|
||||
group_id_int = int(group_id)
|
||||
user_id_int = int(user_id)
|
||||
except (TypeError, ValueError):
|
||||
logger.warning("napcat: invalid group_increase ids group_id={} user_id={}", group_id, user_id)
|
||||
return
|
||||
|
||||
nickname = await self._lookup_member_name(group_id_int, user_id_int)
|
||||
|
||||
# Note: this routes through is_allowed(). For group bots set
|
||||
# `allow_from: ["*"]` (or include the joining user's id) for welcomes
|
||||
# to fire — same trust model as a regular inbound message.
|
||||
await self._handle_message(
|
||||
sender_id=str(user_id),
|
||||
chat_id=f"group:{group_id}",
|
||||
content=f"[group event] new member {nickname} joined group {group_id}",
|
||||
metadata={
|
||||
"is_group": True,
|
||||
"event": "group_increase",
|
||||
},
|
||||
)
|
||||
|
||||
async def _lookup_member_name(self, group_id: int, user_id: int) -> str:
|
||||
"""Lookup group member nickname. Fallback to user id."""
|
||||
try:
|
||||
resp = await self._call_action(
|
||||
"get_group_member_info",
|
||||
{"group_id": group_id, "user_id": user_id, "no_cache": True},
|
||||
)
|
||||
data = resp.get("data", {})
|
||||
# logger.debug("get_group_member_info: {}", resp)
|
||||
return data.get("card") or data.get("nickname") or str(user_id)
|
||||
except Exception as e:
|
||||
logger.warning("napcat: get_group_member_info failed: {}", e)
|
||||
return str(user_id)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Outbound
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
if self._ws is None:
|
||||
logger.warning("napcat: not connected, dropping outbound message")
|
||||
return
|
||||
|
||||
kind, _, target = msg.chat_id.partition(":")
|
||||
if kind not in ("private", "group") or not target:
|
||||
logger.error("napcat: invalid chat_id '{}'", msg.chat_id)
|
||||
return
|
||||
|
||||
segments: list[dict[str, Any]] = []
|
||||
for ref in msg.media or []:
|
||||
if seg := await self._build_image_segment(ref):
|
||||
segments.append(seg)
|
||||
if text := (msg.content or "").strip():
|
||||
segments.append({"type": "text", "data": {"text": text}})
|
||||
if not segments:
|
||||
return
|
||||
|
||||
params: dict[str, Any] = {"message": segments}
|
||||
if kind == "group":
|
||||
params["message_type"] = "group"
|
||||
params["group_id"] = int(target)
|
||||
else:
|
||||
params["message_type"] = "private"
|
||||
params["user_id"] = int(target)
|
||||
|
||||
resp = await self._call_action("send_msg", params)
|
||||
data = resp.get("data") or {}
|
||||
if (mid := data.get("message_id")) is not None:
|
||||
self._bot_outbound_ids.append(int(mid))
|
||||
|
||||
async def _build_image_segment(self, ref: str) -> dict[str, Any] | None:
|
||||
ref = (ref or "").strip()
|
||||
if not ref:
|
||||
return None
|
||||
if ref.startswith(("http://", "https://")):
|
||||
ok, err = validate_url_target(ref)
|
||||
if not ok:
|
||||
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
|
||||
return None
|
||||
return {"type": "image", "data": {"file": ref}}
|
||||
# Local path → base64 so it works even when napcat runs on a
|
||||
# different host/container than vibe-trading.
|
||||
path = Path(os.path.expanduser(ref)).resolve()
|
||||
if not path.is_file():
|
||||
logger.warning("napcat: local image not found: {}", path)
|
||||
return None
|
||||
data = await asyncio.to_thread(path.read_bytes)
|
||||
return {"type": "image", "data": {"file": "base64://" + base64.b64encode(data).decode()}}
|
||||
|
||||
async def _call_action(
|
||||
self,
|
||||
action: str,
|
||||
params: dict[str, Any],
|
||||
timeout: float = _ACTION_TIMEOUT,
|
||||
) -> dict[str, Any]:
|
||||
if self._ws is None:
|
||||
raise RuntimeError("napcat: not connected")
|
||||
echo = uuid.uuid4().hex
|
||||
loop = asyncio.get_running_loop()
|
||||
fut: asyncio.Future[dict[str, Any]] = loop.create_future()
|
||||
self._pending[echo] = fut
|
||||
try:
|
||||
await self._ws.send(
|
||||
json.dumps({"action": action, "params": params, "echo": echo}, ensure_ascii=False)
|
||||
)
|
||||
resp = await asyncio.wait_for(fut, timeout=timeout)
|
||||
status = resp.get("status")
|
||||
retcode = resp.get("retcode")
|
||||
if (status and status != "ok") or (retcode not in (None, 0)):
|
||||
raise RuntimeError(
|
||||
f"napcat: action {action} failed status={status!r} retcode={retcode!r}"
|
||||
)
|
||||
return resp
|
||||
finally:
|
||||
self._pending.pop(echo, None)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Image download
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _download_image(self, info: dict[str, Any]) -> str | None:
|
||||
url = info.get("url")
|
||||
if not isinstance(url, str):
|
||||
return None
|
||||
# logger.debug("napcat: downloading image from {}", url)
|
||||
if self._http is None:
|
||||
return None
|
||||
ok, err = validate_url_target(url)
|
||||
if not ok:
|
||||
logger.warning("napcat: skip image '{}': {}", url, err)
|
||||
return None
|
||||
max_bytes = self.config.max_image_bytes
|
||||
|
||||
# Reject upfront when napcat tells us the size and it's too big.
|
||||
try:
|
||||
declared_size = int(info["file_size"])
|
||||
if declared_size > max_bytes:
|
||||
logger.warning(
|
||||
"napcat: image declared size={} exceeds max_image_bytes={} url={}",
|
||||
declared_size,
|
||||
max_bytes,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
except (TypeError, KeyError):
|
||||
pass
|
||||
|
||||
try:
|
||||
async with self._http.get(url, allow_redirects=False) as resp:
|
||||
if 300 <= resp.status < 400:
|
||||
logger.warning("napcat: image download redirect rejected url={}", url)
|
||||
return None
|
||||
if resp.status >= 400:
|
||||
logger.warning("napcat: image download status={} url={}", resp.status, url)
|
||||
return None
|
||||
# Stream until EOF, capping memory at max_bytes. Don't use
|
||||
# content.read(max_bytes+1) — it returns only what's currently
|
||||
# buffered, which truncates chunked responses mid-image.
|
||||
buf = bytearray()
|
||||
truncated = False
|
||||
async for chunk in resp.content.iter_chunked(64 * 1024):
|
||||
buf.extend(chunk)
|
||||
if len(buf) > max_bytes:
|
||||
truncated = True
|
||||
break
|
||||
if truncated:
|
||||
logger.warning(
|
||||
"napcat: image exceeds max_image_bytes={} url={}", max_bytes, url
|
||||
)
|
||||
return None
|
||||
data = bytes(buf)
|
||||
except Exception as e:
|
||||
logger.warning("napcat: image download error url={} err={}", url, e)
|
||||
return None
|
||||
|
||||
filename_hint = info.get("file")
|
||||
if filename_hint:
|
||||
name = safe_filename(filename_hint)
|
||||
else:
|
||||
name = f"{int(time.time() * 1000)}.jpg"
|
||||
path = self._media_root / name
|
||||
try:
|
||||
await asyncio.to_thread(path.write_bytes, data)
|
||||
except OSError as e:
|
||||
logger.warning("napcat: failed to save image: {}", e)
|
||||
return None
|
||||
return str(path)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Pairing module for DM sender approval."""
|
||||
|
||||
from src.channels.pairing.store import (
|
||||
approve_code,
|
||||
deny_code,
|
||||
format_expiry,
|
||||
format_pairing_reply,
|
||||
generate_code,
|
||||
get_approved,
|
||||
handle_pairing_command,
|
||||
is_approved,
|
||||
list_pending,
|
||||
revoke,
|
||||
)
|
||||
|
||||
# Metadata keys used by channels and commands to tag pairing-related messages.
|
||||
PAIRING_CODE_META_KEY = "_pairing_code"
|
||||
PAIRING_COMMAND_META_KEY = "_pairing_command"
|
||||
|
||||
__all__ = [
|
||||
"PAIRING_CODE_META_KEY",
|
||||
"PAIRING_COMMAND_META_KEY",
|
||||
"approve_code",
|
||||
"deny_code",
|
||||
"format_expiry",
|
||||
"format_pairing_reply",
|
||||
"generate_code",
|
||||
"get_approved",
|
||||
"handle_pairing_command",
|
||||
"is_approved",
|
||||
"list_pending",
|
||||
"revoke",
|
||||
]
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Pairing store for DM sender approval.
|
||||
|
||||
Persistent storage at ``~/.vibe-trading/pairing.json`` keeps approved senders
|
||||
and pending pairing codes per channel. The store is designed for
|
||||
private-assistant scale: small JSON file, simple locking, no external DB.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import secrets
|
||||
import string
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.config.paths import get_data_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# threading.Lock is used so store functions remain callable from both sync CLI
|
||||
# and async channel handlers. At private-assistant scale (small JSON file,
|
||||
# sub-millisecond operations) the brief block is acceptable.
|
||||
_LOCK = threading.Lock()
|
||||
_ALPHABET = string.ascii_uppercase + string.digits
|
||||
_CODE_LENGTH = 8 # e.g. ABCD-EFGH
|
||||
_TTL_DEFAULT_S = 600 # 10 minutes
|
||||
|
||||
|
||||
def _store_path() -> Path:
|
||||
return get_data_dir() / "pairing.json"
|
||||
|
||||
|
||||
def _write_text_atomic(path: Path, text: str) -> None:
|
||||
"""Write *text* to *path* atomically via a temp file + rename."""
|
||||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||||
tmp.write_text(text, encoding="utf-8")
|
||||
tmp.replace(path)
|
||||
|
||||
|
||||
def _load() -> dict[str, Any]:
|
||||
path = _store_path()
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except FileNotFoundError:
|
||||
return {"approved": {}, "pending": {}}
|
||||
except (json.JSONDecodeError, OSError):
|
||||
logger.warning("Corrupted pairing store, resetting")
|
||||
return {"approved": {}, "pending": {}}
|
||||
|
||||
# Convert approved lists to str sets for O(1) lookup.
|
||||
for channel, users in data.get("approved", {}).items():
|
||||
data["approved"][channel] = {str(u) for u in users}
|
||||
return data
|
||||
|
||||
|
||||
def _save(data: dict[str, Any]) -> None:
|
||||
path = _store_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Convert sets back to lists for JSON serialization
|
||||
payload = {
|
||||
"approved": {ch: sorted(list(users)) for ch, users in data.get("approved", {}).items()},
|
||||
"pending": dict(data.get("pending", {})),
|
||||
}
|
||||
_write_text_atomic(path, json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
|
||||
|
||||
def _gc_pending(data: dict[str, Any]) -> None:
|
||||
"""Remove expired pending entries in-place."""
|
||||
now = time.time()
|
||||
pending: dict[str, Any] = data.get("pending", {})
|
||||
expired = [code for code, info in pending.items() if info.get("expires_at", 0) < now]
|
||||
for code in expired:
|
||||
del pending[code]
|
||||
|
||||
|
||||
def generate_code(
|
||||
channel: str,
|
||||
sender_id: str,
|
||||
ttl: int = _TTL_DEFAULT_S,
|
||||
) -> str:
|
||||
"""Create a new pairing code for *sender_id* on *channel*.
|
||||
|
||||
Returns the code (e.g. ``"ABCD-EFGH"``).
|
||||
"""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
_gc_pending(data)
|
||||
raw = "".join(secrets.choice(_ALPHABET) for _ in range(_CODE_LENGTH))
|
||||
code = f"{raw[:4]}-{raw[4:]}"
|
||||
|
||||
data.setdefault("pending", {})[code] = {
|
||||
"channel": channel,
|
||||
"sender_id": str(sender_id),
|
||||
"created_at": time.time(),
|
||||
"expires_at": time.time() + ttl,
|
||||
}
|
||||
_save(data)
|
||||
logger.info("Generated pairing code %s for %s@%s", code, sender_id, channel)
|
||||
return code
|
||||
|
||||
|
||||
def approve_code(code: str) -> tuple[str, str] | None:
|
||||
"""Approve a pending pairing code.
|
||||
|
||||
Returns ``(channel, sender_id)`` on success, or ``None`` if the code
|
||||
does not exist or has expired.
|
||||
"""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
_gc_pending(data)
|
||||
pending: dict[str, Any] = data.get("pending", {})
|
||||
info = pending.pop(code, None)
|
||||
if info is None:
|
||||
return None
|
||||
channel = info["channel"]
|
||||
sender_id = str(info["sender_id"])
|
||||
data.setdefault("approved", {}).setdefault(channel, set()).add(sender_id)
|
||||
_save(data)
|
||||
logger.info("Approved pairing code %s for %s@%s", code, sender_id, channel)
|
||||
return channel, sender_id
|
||||
|
||||
|
||||
def deny_code(code: str) -> bool:
|
||||
"""Reject and discard a pending pairing code.
|
||||
|
||||
Returns ``True`` if the code existed and was removed.
|
||||
"""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
_gc_pending(data)
|
||||
pending: dict[str, Any] = data.get("pending", {})
|
||||
if code in pending:
|
||||
del pending[code]
|
||||
_save(data)
|
||||
logger.info("Denied pairing code %s", code)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_approved(channel: str, sender_id: str) -> bool:
|
||||
"""Check whether *sender_id* has been approved on *channel*."""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
approved: dict[str, set[str]] = data.get("approved", {})
|
||||
return str(sender_id) in approved.get(channel, set())
|
||||
|
||||
|
||||
def list_pending() -> list[dict[str, Any]]:
|
||||
"""Return all non-expired pending pairing requests."""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
_gc_pending(data)
|
||||
return [
|
||||
{"code": code, **info}
|
||||
for code, info in data.get("pending", {}).items()
|
||||
]
|
||||
|
||||
|
||||
def revoke(channel: str, sender_id: str) -> bool:
|
||||
"""Remove an approved sender from *channel*.
|
||||
|
||||
Returns ``True`` if the sender was present and removed.
|
||||
"""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
approved: dict[str, set[str]] = data.get("approved", {})
|
||||
users = approved.get(channel, set())
|
||||
sid = str(sender_id)
|
||||
if sid in users:
|
||||
users.discard(sid)
|
||||
if not users:
|
||||
del approved[channel]
|
||||
_save(data)
|
||||
logger.info("Revoked %s from %s", sid, channel)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_approved(channel: str) -> list[str]:
|
||||
"""Return all approved sender IDs for *channel*."""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
return sorted(data.get("approved", {}).get(channel, set()))
|
||||
|
||||
|
||||
def format_pairing_reply(code: str) -> str:
|
||||
"""Return the pairing-code message sent to unrecognised DM senders."""
|
||||
return (
|
||||
"Hi there! This assistant only responds to approved users.\n\n"
|
||||
f"Your pairing code is: `{code}`\n\n"
|
||||
"To get access, ask the owner to approve this code:\n"
|
||||
f"- In this chat: send `/pairing approve {code}`"
|
||||
)
|
||||
|
||||
|
||||
def format_expiry(expires_at: float) -> str:
|
||||
"""Return a human-readable expiry string (e.g. ``"120s"`` or ``"expired"``)."""
|
||||
remaining = int(expires_at - time.time())
|
||||
return f"{remaining}s" if remaining > 0 else "expired"
|
||||
|
||||
|
||||
def handle_pairing_command(channel: str, subcommand_text: str) -> str:
|
||||
"""Execute a pairing subcommand and return the reply text.
|
||||
|
||||
This is a pure function (no side effects other than store mutations)
|
||||
so it can be used from both the CLI and the agent CommandRouter.
|
||||
"""
|
||||
parts = subcommand_text.split()
|
||||
sub = parts[0] if parts else "list"
|
||||
arg = parts[1] if len(parts) > 1 else None
|
||||
|
||||
if sub in ("list",):
|
||||
pending = list_pending()
|
||||
if not pending:
|
||||
return "No pending pairing requests."
|
||||
lines = ["Pending pairing requests:"]
|
||||
for item in pending:
|
||||
expiry = format_expiry(item.get("expires_at", 0))
|
||||
lines.append(
|
||||
f"- `{item['code']}` | {item['channel']} | {item['sender_id']} | {expiry}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
elif sub == "approve":
|
||||
if arg is None:
|
||||
return "Usage: `/pairing approve <code>`"
|
||||
result = approve_code(arg)
|
||||
if result is None:
|
||||
return f"Invalid or expired pairing code: `{arg}`"
|
||||
ch, sid = result
|
||||
return f"Approved pairing code `{arg}` — {sid} can now access {ch}"
|
||||
|
||||
elif sub == "deny":
|
||||
if arg is None:
|
||||
return "Usage: `/pairing deny <code>`"
|
||||
if deny_code(arg):
|
||||
return f"Denied pairing code `{arg}`"
|
||||
return f"Pairing code `{arg}` not found or already expired"
|
||||
|
||||
elif sub == "revoke":
|
||||
if len(parts) == 2:
|
||||
return (
|
||||
f"Revoked {arg} from {channel}"
|
||||
if revoke(channel, arg)
|
||||
else f"{arg} was not in the approved list for {channel}"
|
||||
)
|
||||
if len(parts) == 3:
|
||||
return (
|
||||
f"Revoked {parts[2]} from {arg}"
|
||||
if revoke(arg, parts[2])
|
||||
else f"{parts[2]} was not in the approved list for {arg}"
|
||||
)
|
||||
return "Usage: `/pairing revoke <user_id>` or `/pairing revoke <channel> <user_id>`"
|
||||
|
||||
return (
|
||||
"Unknown pairing command.\n"
|
||||
"Usage: `/pairing [list|approve <code>|deny <code>|revoke <user_id>|revoke <channel> <user_id>]`"
|
||||
)
|
||||
@@ -0,0 +1,701 @@
|
||||
"""QQ channel implementation using botpy SDK.
|
||||
|
||||
Inbound:
|
||||
- Parse QQ botpy messages (C2C / Group)
|
||||
- Download attachments to media dir using chunked streaming write (memory-safe)
|
||||
- Publish to the Vibe-Trading bus via BaseChannel._handle_message()
|
||||
- Content includes a clear, actionable "Received files:" list with local paths
|
||||
|
||||
Outbound:
|
||||
- Send attachments (msg.media) first via QQ rich media API (base64 upload + msg_type=7)
|
||||
- Then send text (plain or markdown)
|
||||
- msg.media supports local paths, file:// paths, and http(s) URLs
|
||||
|
||||
Notes:
|
||||
- QQ restricts many audio/video formats. We conservatively classify as image vs file.
|
||||
- Attachment structures differ across botpy versions; we try multiple field candidates.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections import deque
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
import aiohttp
|
||||
import logging; logger = logging.getLogger(__name__)
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from pydantic import BaseModel
|
||||
from src.channels.utils import validate_url_target
|
||||
# logging_bridge not needed (using stdlib logging)
|
||||
|
||||
try:
|
||||
from src.channels.utils import get_media_dir
|
||||
except Exception: # pragma: no cover
|
||||
get_media_dir = None # type: ignore
|
||||
|
||||
try:
|
||||
import botpy
|
||||
from botpy.http import Route
|
||||
|
||||
QQ_AVAILABLE = True
|
||||
except ImportError: # pragma: no cover
|
||||
QQ_AVAILABLE = False
|
||||
botpy = None
|
||||
Route = None
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botpy.message import BaseMessage, C2CMessage, GroupMessage
|
||||
from botpy.types.message import Media
|
||||
|
||||
|
||||
# QQ rich media file_type: 1=image, 4=file
|
||||
# (2=voice, 3=video are restricted; we only use image vs file)
|
||||
QQ_FILE_TYPE_IMAGE = 1
|
||||
QQ_FILE_TYPE_FILE = 4
|
||||
|
||||
_IMAGE_EXTS = {
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".gif",
|
||||
".bmp",
|
||||
".webp",
|
||||
".tif",
|
||||
".tiff",
|
||||
".ico",
|
||||
".svg",
|
||||
}
|
||||
|
||||
# Replace unsafe characters with "_", keep Chinese and common safe punctuation.
|
||||
_SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]()【】\u4e00-\u9fff]+", re.UNICODE)
|
||||
|
||||
|
||||
def _sanitize_filename(name: str) -> str:
|
||||
"""Sanitize filename to avoid traversal and problematic chars."""
|
||||
name = (name or "").strip()
|
||||
name = Path(name).name
|
||||
name = _SAFE_NAME_RE.sub("_", name).strip("._ ")
|
||||
return name
|
||||
|
||||
|
||||
def _is_image_name(name: str) -> bool:
|
||||
return Path(name).suffix.lower() in _IMAGE_EXTS
|
||||
|
||||
|
||||
def _guess_send_file_type(filename: str) -> int:
|
||||
"""Conservative send type: images -> 1, else -> 4."""
|
||||
ext = Path(filename).suffix.lower()
|
||||
mime, _ = mimetypes.guess_type(filename)
|
||||
if ext in _IMAGE_EXTS or (mime and mime.startswith("image/")):
|
||||
return QQ_FILE_TYPE_IMAGE
|
||||
return QQ_FILE_TYPE_FILE
|
||||
|
||||
|
||||
def _make_bot_class(channel: QQChannel) -> type[botpy.Client]:
|
||||
"""Create a botpy Client subclass bound to the given channel."""
|
||||
intents = botpy.Intents(public_messages=True, direct_message=True)
|
||||
|
||||
class _Bot(botpy.Client):
|
||||
def __init__(self):
|
||||
# Disable botpy's file log — vibe-trading uses loguru; default "botpy.log" fails on read-only fs
|
||||
super().__init__(intents=intents, ext_handlers=False)
|
||||
|
||||
async def on_ready(self):
|
||||
logger.info("QQ bot ready: {}", self.robot.name)
|
||||
|
||||
async def on_c2c_message_create(self, message: C2CMessage):
|
||||
await channel._on_message(message, is_group=False)
|
||||
|
||||
async def on_group_at_message_create(self, message: GroupMessage):
|
||||
await channel._on_message(message, is_group=True)
|
||||
|
||||
async def on_direct_message_create(self, message):
|
||||
await channel._on_message(message, is_group=False)
|
||||
|
||||
return _Bot
|
||||
|
||||
|
||||
class QQConfig(BaseModel):
|
||||
"""QQ channel configuration using botpy SDK."""
|
||||
|
||||
enabled: bool = False
|
||||
app_id: str = ""
|
||||
secret: str = ""
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
msg_format: Literal["plain", "markdown"] = "plain"
|
||||
ack_message: str = "⏳ Processing..."
|
||||
|
||||
# Optional: directory to save inbound attachments. If empty, use vibe-trading get_media_dir("qq").
|
||||
media_dir: str = ""
|
||||
|
||||
# Download tuning
|
||||
download_chunk_size: int = 1024 * 256 # 256KB
|
||||
download_max_bytes: int = 1024 * 1024 * 200 # 200MB safety limit
|
||||
|
||||
|
||||
class QQChannel(BaseChannel):
|
||||
"""QQ channel using botpy SDK with WebSocket connection."""
|
||||
|
||||
name = "qq"
|
||||
display_name = "QQ"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return QQConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = QQConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: QQConfig = config
|
||||
|
||||
self._client: botpy.Client | None = None
|
||||
self._http: aiohttp.ClientSession | None = None
|
||||
|
||||
self._processed_ids: deque[str] = deque(maxlen=1000)
|
||||
self._msg_seq: int = 1 # used to avoid QQ API dedup
|
||||
self._chat_type_cache: dict[str, str] = {}
|
||||
|
||||
self._media_root: Path = self._init_media_root()
|
||||
|
||||
# ---------------------------
|
||||
# Lifecycle
|
||||
# ---------------------------
|
||||
|
||||
def _init_media_root(self) -> Path:
|
||||
"""Choose a directory for saving inbound attachments."""
|
||||
if self.config.media_dir:
|
||||
root = Path(self.config.media_dir).expanduser()
|
||||
elif get_media_dir:
|
||||
try:
|
||||
root = Path(get_media_dir("qq"))
|
||||
except Exception:
|
||||
root = Path.home() / ".vibe-trading" / "media" / "qq"
|
||||
else:
|
||||
root = Path.home() / ".vibe-trading" / "media" / "qq"
|
||||
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
self.logger.info("media directory: {}", str(root))
|
||||
return root
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the QQ bot with auto-reconnect loop."""
|
||||
# (stdlib logging handles botpy SDK output via propagation)
|
||||
if not QQ_AVAILABLE:
|
||||
self.logger.error("SDK not installed. Run: pip install qq-botpy")
|
||||
return
|
||||
|
||||
if not self.config.app_id or not self.config.secret:
|
||||
self.logger.error("app_id and secret not configured")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))
|
||||
|
||||
self._client = _make_bot_class(self)()
|
||||
self.logger.info("bot started (C2C & Group supported)")
|
||||
await self._run_bot()
|
||||
|
||||
async def _run_bot(self) -> None:
|
||||
"""Run the bot connection with auto-reconnect."""
|
||||
while self._running:
|
||||
try:
|
||||
await self._client.start(appid=self.config.app_id, secret=self.config.secret)
|
||||
except Exception as e:
|
||||
self.logger.warning("bot error: {}", e)
|
||||
if self._running:
|
||||
self.logger.info("Reconnecting bot in 5 seconds...")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop bot and cleanup resources."""
|
||||
self._running = False
|
||||
if self._client:
|
||||
with suppress(Exception):
|
||||
await self._client.close()
|
||||
self._client = None
|
||||
|
||||
if self._http:
|
||||
with suppress(Exception):
|
||||
await self._http.close()
|
||||
self._http = None
|
||||
|
||||
self.logger.info("bot stopped")
|
||||
|
||||
# ---------------------------
|
||||
# Outbound (send)
|
||||
# ---------------------------
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send attachments first, then text."""
|
||||
try:
|
||||
if not self._client:
|
||||
self.logger.warning("client not initialized")
|
||||
return
|
||||
|
||||
msg_id = msg.metadata.get("message_id")
|
||||
chat_type = self._chat_type_cache.get(msg.chat_id, "c2c")
|
||||
is_group = chat_type == "group"
|
||||
|
||||
# 1) Send media
|
||||
for media_ref in msg.media or []:
|
||||
ok = await self._send_media(
|
||||
chat_id=msg.chat_id,
|
||||
media_ref=media_ref,
|
||||
msg_id=msg_id,
|
||||
is_group=is_group,
|
||||
)
|
||||
if not ok:
|
||||
filename = (
|
||||
os.path.basename(urlparse(media_ref).path)
|
||||
or os.path.basename(media_ref)
|
||||
or "file"
|
||||
)
|
||||
await self._send_text_only(
|
||||
chat_id=msg.chat_id,
|
||||
is_group=is_group,
|
||||
msg_id=msg_id,
|
||||
content=f"[Attachment send failed: {filename}]",
|
||||
)
|
||||
|
||||
# 2) Send text
|
||||
if msg.content and msg.content.strip():
|
||||
await self._send_text_only(
|
||||
chat_id=msg.chat_id,
|
||||
is_group=is_group,
|
||||
msg_id=msg_id,
|
||||
content=msg.content.strip(),
|
||||
)
|
||||
except (aiohttp.ClientError, OSError):
|
||||
# Network / transport errors — propagate so ChannelManager can retry
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.exception("Error sending message to chat_id={}", msg.chat_id)
|
||||
|
||||
async def _send_text_only(
|
||||
self,
|
||||
chat_id: str,
|
||||
is_group: bool,
|
||||
msg_id: str | None,
|
||||
content: str,
|
||||
) -> None:
|
||||
"""Send a plain/markdown text message."""
|
||||
if not self._client:
|
||||
return
|
||||
|
||||
self._msg_seq += 1
|
||||
use_markdown = self.config.msg_format == "markdown"
|
||||
payload: dict[str, Any] = {
|
||||
"msg_type": 2 if use_markdown else 0,
|
||||
"msg_id": msg_id,
|
||||
"msg_seq": self._msg_seq,
|
||||
}
|
||||
if use_markdown:
|
||||
payload["markdown"] = {"content": content}
|
||||
else:
|
||||
payload["content"] = content
|
||||
|
||||
if is_group:
|
||||
await self._client.api.post_group_message(group_openid=chat_id, **payload)
|
||||
else:
|
||||
await self._client.api.post_c2c_message(openid=chat_id, **payload)
|
||||
|
||||
async def _send_media(
|
||||
self,
|
||||
chat_id: str,
|
||||
media_ref: str,
|
||||
msg_id: str | None,
|
||||
is_group: bool,
|
||||
) -> bool:
|
||||
"""Read bytes -> base64 upload -> msg_type=7 send."""
|
||||
if not self._client:
|
||||
return False
|
||||
|
||||
data, filename = await self._read_media_bytes(media_ref)
|
||||
if not data or not filename:
|
||||
return False
|
||||
|
||||
try:
|
||||
file_type = _guess_send_file_type(filename)
|
||||
file_data_b64 = base64.b64encode(data).decode()
|
||||
|
||||
media_obj = await self._post_base64file(
|
||||
chat_id=chat_id,
|
||||
is_group=is_group,
|
||||
file_type=file_type,
|
||||
file_data=file_data_b64,
|
||||
file_name=filename,
|
||||
srv_send_msg=False,
|
||||
)
|
||||
if not media_obj:
|
||||
self.logger.error("media upload failed: empty response")
|
||||
return False
|
||||
|
||||
self._msg_seq += 1
|
||||
if is_group:
|
||||
await self._client.api.post_group_message(
|
||||
group_openid=chat_id,
|
||||
msg_type=7,
|
||||
msg_id=msg_id,
|
||||
msg_seq=self._msg_seq,
|
||||
media=media_obj,
|
||||
)
|
||||
else:
|
||||
await self._client.api.post_c2c_message(
|
||||
openid=chat_id,
|
||||
msg_type=7,
|
||||
msg_id=msg_id,
|
||||
msg_seq=self._msg_seq,
|
||||
media=media_obj,
|
||||
)
|
||||
|
||||
self.logger.info("media sent: {}", filename)
|
||||
return True
|
||||
except (aiohttp.ClientError, OSError) as e:
|
||||
# Network / transport errors — propagate for retry by caller
|
||||
self.logger.warning("send media network error filename={} err={}", filename, e)
|
||||
raise
|
||||
except Exception:
|
||||
# API-level or other non-network errors — return False so send() can fallback
|
||||
self.logger.exception("send media failed filename={}", filename)
|
||||
return False
|
||||
|
||||
async def _read_media_bytes(self, media_ref: str) -> tuple[bytes | None, str | None]:
|
||||
"""Read bytes from http(s) or local file path; return (data, filename)."""
|
||||
media_ref = (media_ref or "").strip()
|
||||
if not media_ref:
|
||||
return None, None
|
||||
|
||||
# Local file: plain path or file:// URI
|
||||
if not media_ref.startswith("http://") and not media_ref.startswith("https://"):
|
||||
try:
|
||||
if media_ref.startswith("file://"):
|
||||
parsed = urlparse(media_ref)
|
||||
# Windows: path in netloc; Unix: path in path
|
||||
raw = parsed.path or parsed.netloc
|
||||
local_path = Path(unquote(raw))
|
||||
else:
|
||||
local_path = Path(os.path.expanduser(media_ref))
|
||||
|
||||
if not local_path.is_file():
|
||||
self.logger.warning("outbound media file not found: {}", str(local_path))
|
||||
return None, None
|
||||
|
||||
data = await asyncio.to_thread(local_path.read_bytes)
|
||||
return data, local_path.name
|
||||
except Exception as e:
|
||||
self.logger.warning("outbound media read error ref={} err={}", media_ref, e)
|
||||
return None, None
|
||||
|
||||
# Remote URL
|
||||
ok, err = validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
|
||||
return None, None
|
||||
|
||||
if not self._http:
|
||||
self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))
|
||||
try:
|
||||
async with self._http.get(media_ref, allow_redirects=True) as resp:
|
||||
if resp.status >= 400:
|
||||
self.logger.warning(
|
||||
"outbound media download failed status={} url={}",
|
||||
resp.status,
|
||||
media_ref,
|
||||
)
|
||||
return None, None
|
||||
data = await resp.read()
|
||||
if not data:
|
||||
return None, None
|
||||
filename = os.path.basename(urlparse(media_ref).path) or "file.bin"
|
||||
return data, filename
|
||||
except Exception as e:
|
||||
self.logger.warning("outbound media download error url={} err={}", media_ref, e)
|
||||
return None, None
|
||||
|
||||
# https://github.com/tencent-connect/botpy/issues/198
|
||||
# https://bot.q.qq.com/wiki/develop/api-v2/server-inter/message/send-receive/rich-media.html
|
||||
async def _post_base64file(
|
||||
self,
|
||||
chat_id: str,
|
||||
is_group: bool,
|
||||
file_type: int,
|
||||
file_data: str,
|
||||
file_name: str | None = None,
|
||||
srv_send_msg: bool = False,
|
||||
) -> Media:
|
||||
"""Upload base64-encoded file and return Media object."""
|
||||
if not self._client:
|
||||
raise RuntimeError("QQ client not initialized")
|
||||
|
||||
if is_group:
|
||||
endpoint = "/v2/groups/{group_openid}/files"
|
||||
id_key = "group_openid"
|
||||
else:
|
||||
endpoint = "/v2/users/{openid}/files"
|
||||
id_key = "openid"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
id_key: chat_id,
|
||||
"file_type": file_type,
|
||||
"file_data": file_data,
|
||||
"srv_send_msg": srv_send_msg,
|
||||
}
|
||||
# Only pass file_name for non-image types (file_type=4).
|
||||
# Passing file_name for images causes QQ client to render them as
|
||||
# file attachments instead of inline images.
|
||||
if file_type != QQ_FILE_TYPE_IMAGE and file_name:
|
||||
payload["file_name"] = file_name
|
||||
|
||||
route = Route("POST", endpoint, **{id_key: chat_id})
|
||||
result = await self._client.api._http.request(route, json=payload)
|
||||
|
||||
# Extract only the file_info field to avoid extra fields (file_uuid, ttl, etc.)
|
||||
# that may confuse QQ client when sending the media object.
|
||||
if isinstance(result, dict) and "file_info" in result:
|
||||
return {"file_info": result["file_info"]}
|
||||
return result
|
||||
|
||||
# ---------------------------
|
||||
# Inbound (receive)
|
||||
# ---------------------------
|
||||
|
||||
async def _on_message(self, data: C2CMessage | GroupMessage, is_group: bool = False) -> None:
|
||||
"""Parse inbound message, download attachments, and publish to the bus."""
|
||||
try:
|
||||
if is_group:
|
||||
chat_id = data.group_openid
|
||||
user_id = data.author.member_openid
|
||||
chat_type = "group"
|
||||
else:
|
||||
chat_id = str(
|
||||
getattr(data.author, "id", None)
|
||||
or getattr(data.author, "user_openid", "unknown")
|
||||
)
|
||||
user_id = chat_id
|
||||
chat_type = "c2c"
|
||||
|
||||
content = (data.content or "").strip()
|
||||
|
||||
if data.id in self._processed_ids:
|
||||
return
|
||||
self._processed_ids.append(data.id)
|
||||
self._chat_type_cache[chat_id] = chat_type
|
||||
|
||||
# Early permission check — avoid attachment downloads and ack side effects
|
||||
# for unauthorized users. C2C messages can receive pairing codes;
|
||||
# group messages remain silently ignored.
|
||||
if not self.is_allowed(user_id):
|
||||
if not is_group:
|
||||
await self._handle_message(
|
||||
sender_id=user_id,
|
||||
chat_id=chat_id,
|
||||
content="",
|
||||
is_dm=True,
|
||||
)
|
||||
return
|
||||
|
||||
# the data used by tests don't contain attachments property
|
||||
# so we use getattr with a default of [] to avoid AttributeError in tests
|
||||
attachments = getattr(data, "attachments", None) or []
|
||||
media_paths, recv_lines, att_meta = await self._handle_attachments(attachments)
|
||||
|
||||
# Compose content that always contains actionable saved paths
|
||||
if recv_lines:
|
||||
tag = (
|
||||
"[Image]"
|
||||
if any(_is_image_name(Path(p).name) for p in media_paths)
|
||||
else "[File]"
|
||||
)
|
||||
file_block = "Received files:\n" + "\n".join(recv_lines)
|
||||
content = (
|
||||
f"{content}\n\n{file_block}".strip() if content else f"{tag}\n{file_block}"
|
||||
)
|
||||
|
||||
if not content and not media_paths:
|
||||
return
|
||||
|
||||
if self.config.ack_message:
|
||||
try:
|
||||
await self._send_text_only(
|
||||
chat_id=chat_id,
|
||||
is_group=is_group,
|
||||
msg_id=data.id,
|
||||
content=self.config.ack_message,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.debug("ack message failed for chat_id={}", chat_id)
|
||||
|
||||
await self._handle_message(
|
||||
sender_id=user_id,
|
||||
chat_id=chat_id,
|
||||
content=content,
|
||||
media=media_paths if media_paths else None,
|
||||
metadata={
|
||||
"message_id": data.id,
|
||||
"attachments": att_meta,
|
||||
},
|
||||
is_dm=not is_group,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Error handling inbound message id={}", getattr(data, "id", "?"))
|
||||
|
||||
async def _handle_attachments(
|
||||
self,
|
||||
attachments: list[BaseMessage._Attachments],
|
||||
) -> tuple[list[str], list[str], list[dict[str, Any]]]:
|
||||
"""Extract, download (chunked), and format attachments for agent consumption."""
|
||||
media_paths: list[str] = []
|
||||
recv_lines: list[str] = []
|
||||
att_meta: list[dict[str, Any]] = []
|
||||
|
||||
if not attachments:
|
||||
return media_paths, recv_lines, att_meta
|
||||
|
||||
for att in attachments:
|
||||
url = getattr(att, "url", None) or ""
|
||||
filename = getattr(att, "filename", None) or ""
|
||||
ctype = getattr(att, "content_type", None) or ""
|
||||
|
||||
self.logger.info("Downloading file: {}", filename or url)
|
||||
local_path = await self._download_to_media_dir_chunked(url, filename_hint=filename)
|
||||
|
||||
att_meta.append(
|
||||
{
|
||||
"url": url,
|
||||
"filename": filename,
|
||||
"content_type": ctype,
|
||||
"saved_path": local_path,
|
||||
}
|
||||
)
|
||||
|
||||
if local_path:
|
||||
media_paths.append(local_path)
|
||||
shown_name = filename or os.path.basename(local_path)
|
||||
recv_lines.append(f"- {shown_name}\n saved: {local_path}")
|
||||
else:
|
||||
shown_name = filename or url
|
||||
recv_lines.append(f"- {shown_name}\n saved: [download failed]")
|
||||
|
||||
return media_paths, recv_lines, att_meta
|
||||
|
||||
async def _download_to_media_dir_chunked(
|
||||
self,
|
||||
url: str,
|
||||
filename_hint: str = "",
|
||||
) -> str | None:
|
||||
"""Download an inbound attachment using streaming chunk write.
|
||||
|
||||
Uses chunked streaming to avoid loading large files into memory.
|
||||
Enforces a max download size and writes to a .part temp file
|
||||
that is atomically renamed on success.
|
||||
"""
|
||||
# Handle protocol-relative URLs (e.g. "//multimedia.nt.qq.com/...")
|
||||
if url.startswith("//"):
|
||||
url = f"https:{url}"
|
||||
|
||||
if not self._http:
|
||||
self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))
|
||||
|
||||
safe = _sanitize_filename(filename_hint)
|
||||
ts = int(time.time() * 1000)
|
||||
tmp_path: Path | None = None
|
||||
|
||||
try:
|
||||
async with self._http.get(
|
||||
url,
|
||||
timeout=aiohttp.ClientTimeout(total=120),
|
||||
allow_redirects=True,
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
self.logger.warning("download failed: status={} url={}", resp.status, url)
|
||||
return None
|
||||
|
||||
ctype = (resp.headers.get("Content-Type") or "").lower()
|
||||
|
||||
# Infer extension: url -> filename_hint -> content-type -> fallback
|
||||
ext = Path(urlparse(url).path).suffix
|
||||
if not ext:
|
||||
ext = Path(filename_hint).suffix
|
||||
if not ext:
|
||||
if "png" in ctype:
|
||||
ext = ".png"
|
||||
elif "jpeg" in ctype or "jpg" in ctype:
|
||||
ext = ".jpg"
|
||||
elif "gif" in ctype:
|
||||
ext = ".gif"
|
||||
elif "webp" in ctype:
|
||||
ext = ".webp"
|
||||
elif "pdf" in ctype:
|
||||
ext = ".pdf"
|
||||
else:
|
||||
ext = ".bin"
|
||||
|
||||
if safe:
|
||||
if not Path(safe).suffix:
|
||||
safe = safe + ext
|
||||
filename = safe
|
||||
else:
|
||||
filename = f"qq_file_{ts}{ext}"
|
||||
|
||||
target = self._media_root / filename
|
||||
if target.exists():
|
||||
target = self._media_root / f"{target.stem}_{ts}{target.suffix}"
|
||||
|
||||
tmp_path = target.with_suffix(target.suffix + ".part")
|
||||
|
||||
# Stream write
|
||||
downloaded = 0
|
||||
chunk_size = max(1024, int(self.config.download_chunk_size or 262144))
|
||||
max_bytes = max(
|
||||
1024 * 1024, int(self.config.download_max_bytes or (200 * 1024 * 1024))
|
||||
)
|
||||
|
||||
def _open_tmp():
|
||||
tmp_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return open(tmp_path, "wb") # noqa: SIM115
|
||||
|
||||
f = await asyncio.to_thread(_open_tmp)
|
||||
try:
|
||||
async for chunk in resp.content.iter_chunked(chunk_size):
|
||||
if not chunk:
|
||||
continue
|
||||
downloaded += len(chunk)
|
||||
if downloaded > max_bytes:
|
||||
self.logger.warning(
|
||||
"download exceeded max_bytes={} url={} -> abort",
|
||||
max_bytes,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
await asyncio.to_thread(f.write, chunk)
|
||||
finally:
|
||||
await asyncio.to_thread(f.close)
|
||||
|
||||
# Atomic rename
|
||||
await asyncio.to_thread(os.replace, tmp_path, target)
|
||||
tmp_path = None # mark as moved
|
||||
self.logger.info("file saved: {}", str(target))
|
||||
return str(target)
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("download error")
|
||||
return None
|
||||
finally:
|
||||
# Cleanup partial file
|
||||
if tmp_path is not None:
|
||||
with suppress(Exception):
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Auto-discovery for built-in channel modules and external plugins."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import logging
|
||||
import pkgutil
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.channels.base import BaseChannel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_INTERNAL = frozenset({"base", "bus", "config", "manager", "pairing", "registry", "runtime", "utils"})
|
||||
|
||||
_INSTALL_HINTS: dict[str, str] = {
|
||||
"dingtalk": "pip install 'vibe-trading-ai[dingtalk]'",
|
||||
"discord": "pip install 'vibe-trading-ai[discord]'",
|
||||
"email": "No extra Python package required; configure channels.email in the agent config.",
|
||||
"feishu": "pip install 'vibe-trading-ai[feishu]'",
|
||||
"matrix": "pip install 'vibe-trading-ai[matrix]'",
|
||||
"mochat": "pip install 'vibe-trading-ai[mochat]'",
|
||||
"msteams": "pip install 'vibe-trading-ai[msteams]'",
|
||||
"napcat": "pip install 'vibe-trading-ai[napcat]'",
|
||||
"qq": "pip install 'vibe-trading-ai[qq]'",
|
||||
"signal": "No extra Python package required; install and run signal-cli-rest-api separately.",
|
||||
"slack": "pip install 'vibe-trading-ai[slack]'",
|
||||
"telegram": "pip install 'vibe-trading-ai[telegram]'",
|
||||
"wecom": "pip install 'vibe-trading-ai[wecom]'",
|
||||
"weixin": "No extra Python package required; configure channels.weixin in the agent config.",
|
||||
"whatsapp": "pip install 'vibe-trading-ai[whatsapp]'",
|
||||
"websocket": "pip install 'vibe-trading-ai[channels]'",
|
||||
}
|
||||
|
||||
_AVAILABILITY_FLAGS: dict[str, tuple[str, ...]] = {
|
||||
"dingtalk": ("DINGTALK_AVAILABLE",),
|
||||
"discord": ("DISCORD_AVAILABLE",),
|
||||
"feishu": ("FEISHU_AVAILABLE",),
|
||||
"msteams": ("MSTEAMS_AVAILABLE",),
|
||||
"qq": ("QQ_AVAILABLE",),
|
||||
"wecom": ("WECOM_AVAILABLE",),
|
||||
}
|
||||
|
||||
_LAZY_IMPORT_PACKAGES: dict[str, tuple[str, ...]] = {
|
||||
"whatsapp": ("neonize",),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChannelAvailability:
|
||||
"""Import-time availability for a channel adapter."""
|
||||
|
||||
name: str
|
||||
available: bool
|
||||
display_name: str
|
||||
error: str = ""
|
||||
install_hint: str = ""
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Return a JSON-serializable availability payload."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"available": self.available,
|
||||
"display_name": self.display_name,
|
||||
"error": self.error,
|
||||
"install_hint": self.install_hint,
|
||||
}
|
||||
|
||||
|
||||
def discover_channel_names() -> list[str]:
|
||||
"""Return all built-in channel module names by scanning the package (zero imports)."""
|
||||
import src.channels as pkg
|
||||
|
||||
return [
|
||||
name
|
||||
for _, name, ispkg in pkgutil.iter_modules(pkg.__path__)
|
||||
if name not in _INTERNAL and not ispkg
|
||||
]
|
||||
|
||||
|
||||
def _channel_class_from_module(module: ModuleType, module_name: str) -> type[BaseChannel]:
|
||||
"""Return the first BaseChannel subclass defined in *module*."""
|
||||
from src.channels.base import BaseChannel as _Base
|
||||
|
||||
for attr in dir(module):
|
||||
obj = getattr(module, attr)
|
||||
if isinstance(obj, type) and issubclass(obj, _Base) and obj is not _Base:
|
||||
return obj
|
||||
raise ImportError(f"No BaseChannel subclass in src.channels.{module_name}")
|
||||
|
||||
|
||||
def _missing_optional_dependency(name: str, module: ModuleType) -> str:
|
||||
"""Return a dependency status error for lazily imported adapter SDKs."""
|
||||
for flag in _AVAILABILITY_FLAGS.get(name, ()):
|
||||
if getattr(module, flag, True) is False:
|
||||
return f"missing optional dependency for {name}"
|
||||
missing_packages = [
|
||||
package
|
||||
for package in _LAZY_IMPORT_PACKAGES.get(name, ())
|
||||
if importlib.util.find_spec(package) is None
|
||||
]
|
||||
if missing_packages:
|
||||
return "missing optional dependency: " + ", ".join(missing_packages)
|
||||
return ""
|
||||
|
||||
|
||||
def load_channel_class(module_name: str) -> type[BaseChannel]:
|
||||
"""Import *module_name* and return the first BaseChannel subclass found."""
|
||||
mod = importlib.import_module(f"src.channels.{module_name}")
|
||||
return _channel_class_from_module(mod, module_name)
|
||||
|
||||
|
||||
def inspect_channel(name: str) -> ChannelAvailability:
|
||||
"""Inspect one channel module without raising dependency errors.
|
||||
|
||||
Args:
|
||||
name: Built-in or plugin channel name.
|
||||
|
||||
Returns:
|
||||
Availability metadata that surfaces install guidance for optional SDKs.
|
||||
"""
|
||||
try:
|
||||
mod = importlib.import_module(f"src.channels.{name}")
|
||||
cls = _channel_class_from_module(mod, name)
|
||||
display = getattr(cls, "display_name", name)
|
||||
missing = _missing_optional_dependency(name, mod)
|
||||
if missing:
|
||||
return ChannelAvailability(
|
||||
name=name,
|
||||
available=False,
|
||||
display_name=str(display),
|
||||
error=missing,
|
||||
install_hint=_INSTALL_HINTS.get(name, f"pip install 'vibe-trading-ai[{name}]'"),
|
||||
)
|
||||
return ChannelAvailability(name=name, available=True, display_name=str(display))
|
||||
except Exception as exc: # noqa: BLE001 - status API must report every adapter
|
||||
return ChannelAvailability(
|
||||
name=name,
|
||||
available=False,
|
||||
display_name=name.replace("_", " ").title(),
|
||||
error=f"{type(exc).__name__}: {exc}",
|
||||
install_hint=_INSTALL_HINTS.get(name, f"pip install 'vibe-trading-ai[{name}]'"),
|
||||
)
|
||||
|
||||
|
||||
def _section_enabled(section: Any) -> bool:
|
||||
if section is None:
|
||||
return False
|
||||
if isinstance(section, Mapping):
|
||||
return bool(section.get("enabled", False))
|
||||
return bool(getattr(section, "enabled", False))
|
||||
|
||||
|
||||
def _config_section(config: Any, name: str) -> Any:
|
||||
if isinstance(config, Mapping):
|
||||
return config.get(name)
|
||||
return getattr(config, name, None)
|
||||
|
||||
|
||||
def inspect_channels(config: Any | None = None) -> dict[str, dict[str, Any]]:
|
||||
"""Inspect all built-in channels and annotate configured/enabled state.
|
||||
|
||||
Args:
|
||||
config: Optional channels config map or model. Extra keys are included
|
||||
so external plugin configs are visible in status output.
|
||||
|
||||
Returns:
|
||||
Mapping from channel name to JSON-serializable status metadata.
|
||||
"""
|
||||
names = set(discover_channel_names())
|
||||
if isinstance(config, Mapping):
|
||||
names.update(str(key) for key in config.keys())
|
||||
elif config is not None:
|
||||
names.update(
|
||||
key
|
||||
for key in dir(config)
|
||||
if not key.startswith("_") and key not in {"model_config", "model_fields"}
|
||||
)
|
||||
|
||||
statuses: dict[str, dict[str, Any]] = {}
|
||||
for name in sorted(names):
|
||||
section = _config_section(config, name) if config is not None else None
|
||||
availability = inspect_channel(name).to_dict()
|
||||
availability.update(
|
||||
{
|
||||
"configured": section is not None,
|
||||
"enabled": _section_enabled(section),
|
||||
"loaded": False,
|
||||
"running": False,
|
||||
}
|
||||
)
|
||||
statuses[name] = availability
|
||||
return statuses
|
||||
|
||||
|
||||
def discover_plugins(
|
||||
enabled_names: set[str] | None = None,
|
||||
) -> dict[str, type[BaseChannel]]:
|
||||
"""Discover external channel plugins registered via entry_points."""
|
||||
from importlib.metadata import entry_points
|
||||
|
||||
plugins: dict[str, type[BaseChannel]] = {}
|
||||
for ep in entry_points(group="vibe_trading.channels"):
|
||||
if enabled_names is not None and ep.name not in enabled_names:
|
||||
continue
|
||||
try:
|
||||
cls = ep.load()
|
||||
plugins[ep.name] = cls
|
||||
except Exception:
|
||||
logger.warning("Failed to load channel plugin '%s': %s", ep.name, exc_info=True)
|
||||
return plugins
|
||||
|
||||
|
||||
def discover_enabled(
|
||||
enabled_names: set[str],
|
||||
*,
|
||||
_names: list[str] | None = None,
|
||||
_include_all_external: bool = False,
|
||||
) -> dict[str, type[BaseChannel]]:
|
||||
"""Return channels whose module names are in *enabled_names*.
|
||||
|
||||
Uses cheap ``pkgutil.iter_modules`` to list names, then imports only
|
||||
those that match — skipping the heavy third-party SDK imports of
|
||||
unneeded channels.
|
||||
"""
|
||||
names = _names if _names is not None else discover_channel_names()
|
||||
result: dict[str, type[BaseChannel]] = {}
|
||||
for modname in names:
|
||||
if modname not in enabled_names:
|
||||
continue
|
||||
try:
|
||||
result[modname] = load_channel_class(modname)
|
||||
except ImportError:
|
||||
logger.debug("Skipping built-in channel '%s': %s", modname, exc_info=True)
|
||||
|
||||
external = discover_plugins(None if _include_all_external else enabled_names)
|
||||
shadowed = set(external) & set(result)
|
||||
if shadowed:
|
||||
logger.warning("Plugin(s) shadowed by built-in channels (ignored): %s", shadowed)
|
||||
if _include_all_external:
|
||||
result.update({k: v for k, v in external.items() if k not in shadowed})
|
||||
else:
|
||||
result.update(
|
||||
{k: v for k, v in external.items() if k not in shadowed and k in enabled_names}
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def discover_all() -> dict[str, type[BaseChannel]]:
|
||||
"""Return all channels: built-in (pkgutil) merged with external (entry_points).
|
||||
|
||||
Built-in channels take priority — an external plugin cannot shadow a built-in name.
|
||||
"""
|
||||
names = discover_channel_names()
|
||||
return discover_enabled(set(names), _names=names, _include_all_external=True)
|
||||
@@ -0,0 +1,215 @@
|
||||
"""IM channel runtime that connects MessageBus traffic to SessionService."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.manager import ChannelManager
|
||||
from src.channels.pairing import PAIRING_COMMAND_META_KEY, handle_pairing_command
|
||||
from src.config.paths import get_data_dir
|
||||
from src.session.models import Message, Session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChannelRuntimeConfig:
|
||||
"""Runtime controls for IM channel processing."""
|
||||
|
||||
reply_timeout_s: float = 600.0
|
||||
poll_interval_s: float = 0.25
|
||||
|
||||
|
||||
class ChannelRuntime:
|
||||
"""Route inbound channel messages into Vibe-Trading sessions."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
bus: MessageBus,
|
||||
session_service: Any,
|
||||
manager: ChannelManager | None,
|
||||
session_map_path: Path | None = None,
|
||||
reply_timeout_s: float = 600.0,
|
||||
poll_interval_s: float = 0.25,
|
||||
) -> None:
|
||||
self.bus = bus
|
||||
self.session_service = session_service
|
||||
self.manager = manager
|
||||
self.config = ChannelRuntimeConfig(
|
||||
reply_timeout_s=reply_timeout_s,
|
||||
poll_interval_s=poll_interval_s,
|
||||
)
|
||||
self.session_map_path = session_map_path or (get_data_dir() / "channels" / "sessions.json")
|
||||
self._session_map: dict[str, str] = {}
|
||||
self._consumer_task: asyncio.Task[None] | None = None
|
||||
self._manager_task: asyncio.Task[Any] | None = None
|
||||
self._handler_tasks: set[asyncio.Task[None]] = set()
|
||||
self._running = False
|
||||
|
||||
async def start(self, *, start_manager: bool = True) -> None:
|
||||
"""Start channel processing and, optionally, platform adapters."""
|
||||
if self._running:
|
||||
return
|
||||
self._session_map = self._load_session_map()
|
||||
self._running = True
|
||||
if start_manager and self.manager is not None:
|
||||
self._manager_task = asyncio.create_task(self.manager.start_all())
|
||||
await asyncio.sleep(0)
|
||||
self._consumer_task = asyncio.create_task(self._consume_loop())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop channel processing and platform adapters."""
|
||||
self._running = False
|
||||
if self._consumer_task is not None:
|
||||
self._consumer_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._consumer_task
|
||||
self._consumer_task = None
|
||||
for task in list(self._handler_tasks):
|
||||
task.cancel()
|
||||
for task in list(self._handler_tasks):
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
self._handler_tasks.clear()
|
||||
if self.manager is not None:
|
||||
await self.manager.stop_all()
|
||||
if self._manager_task is not None:
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._manager_task
|
||||
self._manager_task = None
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
"""Return runtime and channel status."""
|
||||
return {
|
||||
"running": self._running,
|
||||
"inbound_queue": self.bus.inbound_size,
|
||||
"outbound_queue": self.bus.outbound_size,
|
||||
"session_count": len(self._session_map),
|
||||
"channels": self.manager.get_status() if self.manager is not None else {},
|
||||
}
|
||||
|
||||
async def _consume_loop(self) -> None:
|
||||
while True:
|
||||
msg = await self.bus.consume_inbound()
|
||||
task = asyncio.create_task(self._handle_inbound(msg))
|
||||
self._handler_tasks.add(task)
|
||||
task.add_done_callback(self._handler_tasks.discard)
|
||||
|
||||
async def _handle_inbound(self, msg: InboundMessage) -> None:
|
||||
try:
|
||||
if self._is_pairing_command(msg.content):
|
||||
reply = handle_pairing_command(msg.channel, msg.content.split(None, 1)[1] if " " in msg.content else "list")
|
||||
await self.bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel=msg.channel,
|
||||
chat_id=msg.chat_id,
|
||||
content=reply,
|
||||
metadata={PAIRING_COMMAND_META_KEY: True},
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
session_id = self._session_for(msg)
|
||||
result = await self.session_service.send_message(
|
||||
session_id,
|
||||
msg.content,
|
||||
include_shell_tools=False,
|
||||
)
|
||||
attempt_id = result.get("attempt_id") if isinstance(result, dict) else None
|
||||
reply = await self._wait_for_reply(session_id, attempt_id)
|
||||
await self.bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel=msg.channel,
|
||||
chat_id=msg.chat_id,
|
||||
content=reply.content,
|
||||
metadata={
|
||||
"_channel_runtime": True,
|
||||
"attempt_id": attempt_id,
|
||||
"session_id": session_id,
|
||||
},
|
||||
)
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - channel errors must surface to users
|
||||
logger.exception("Channel runtime failed for %s:%s", msg.channel, msg.chat_id)
|
||||
await self.bus.publish_outbound(
|
||||
OutboundMessage(
|
||||
channel=msg.channel,
|
||||
chat_id=msg.chat_id,
|
||||
content=f"Channel runtime error: {type(exc).__name__}: {exc}",
|
||||
metadata={"_channel_runtime": True, "error": True},
|
||||
)
|
||||
)
|
||||
|
||||
def _session_for(self, msg: InboundMessage) -> str:
|
||||
key = msg.session_key
|
||||
existing = self._session_map.get(key)
|
||||
if existing:
|
||||
return existing
|
||||
session = self.session_service.create_session(
|
||||
title=f"{msg.channel}:{msg.chat_id}",
|
||||
config={"channel": msg.channel, "channel_chat_id": msg.chat_id},
|
||||
)
|
||||
session_id = _session_id(session)
|
||||
self._session_map[key] = session_id
|
||||
self._save_session_map()
|
||||
return session_id
|
||||
|
||||
async def _wait_for_reply(self, session_id: str, attempt_id: str | None) -> Message:
|
||||
deadline = time.monotonic() + self.config.reply_timeout_s
|
||||
last_assistant: Message | None = None
|
||||
while time.monotonic() < deadline:
|
||||
messages = self.session_service.get_messages(session_id, limit=200)
|
||||
for message in reversed(messages):
|
||||
if message.role != "assistant":
|
||||
continue
|
||||
if attempt_id and message.linked_attempt_id != attempt_id:
|
||||
if last_assistant is None:
|
||||
last_assistant = message
|
||||
continue
|
||||
return message
|
||||
await asyncio.sleep(self.config.poll_interval_s)
|
||||
if last_assistant is not None:
|
||||
return last_assistant
|
||||
raise TimeoutError("timed out waiting for assistant reply")
|
||||
|
||||
def _load_session_map(self) -> dict[str, str]:
|
||||
try:
|
||||
data = json.loads(self.session_map_path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.warning("Ignoring invalid channel session map at %s", self.session_map_path)
|
||||
return {}
|
||||
if not isinstance(data, dict):
|
||||
return {}
|
||||
return {str(key): str(value) for key, value in data.items() if value}
|
||||
|
||||
def _save_session_map(self) -> None:
|
||||
self.session_map_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = self.session_map_path.with_suffix(self.session_map_path.suffix + ".tmp")
|
||||
tmp.write_text(json.dumps(self._session_map, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
tmp.replace(self.session_map_path)
|
||||
|
||||
@staticmethod
|
||||
def _is_pairing_command(content: str) -> bool:
|
||||
return content.strip() == "/pairing" or content.strip().startswith("/pairing ")
|
||||
|
||||
|
||||
def _session_id(session: Session | dict[str, Any] | Any) -> str:
|
||||
if isinstance(session, Session):
|
||||
return session.session_id
|
||||
if isinstance(session, dict):
|
||||
return str(session["session_id"])
|
||||
return str(getattr(session, "session_id"))
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,740 @@
|
||||
"""Slack channel implementation using Socket Mode."""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from pydantic import Field
|
||||
from slack_sdk.socket_mode.request import SocketModeRequest
|
||||
from slack_sdk.socket_mode.response import SocketModeResponse
|
||||
from slack_sdk.socket_mode.websockets import SocketModeClient
|
||||
from slack_sdk.web.async_client import AsyncWebClient
|
||||
from slackify_markdown import slackify_markdown
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_media_dir
|
||||
from pydantic import BaseModel
|
||||
from src.channels.pairing import is_approved
|
||||
from src.channels.utils import safe_filename, split_message
|
||||
|
||||
|
||||
class SlackDMConfig(BaseModel):
|
||||
"""Slack DM policy configuration."""
|
||||
|
||||
enabled: bool = True
|
||||
policy: str = "open"
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class SlackConfig(BaseModel):
|
||||
"""Slack channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
mode: str = "socket"
|
||||
webhook_path: str = "/slack/events"
|
||||
bot_token: str = ""
|
||||
app_token: str = ""
|
||||
user_token_read_only: bool = True
|
||||
reply_in_thread: bool = True
|
||||
react_emoji: str = "eyes"
|
||||
done_emoji: str = "white_check_mark"
|
||||
include_thread_context: bool = True
|
||||
thread_context_limit: int = 20
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
group_policy: str = "mention"
|
||||
group_allow_from: list[str] = Field(default_factory=list)
|
||||
# When group_policy is "allowlist", also require the bot to be @mentioned
|
||||
# before responding (so it only replies to mentions in approved channels,
|
||||
# instead of every message). No effect for "mention"/"open" policies.
|
||||
group_require_mention: bool = False
|
||||
dm: SlackDMConfig = Field(default_factory=SlackDMConfig)
|
||||
|
||||
|
||||
SLACK_MAX_MESSAGE_LEN = 39_000 # Slack API allows ~40k; leave margin
|
||||
SLACK_DOWNLOAD_TIMEOUT = 30.0
|
||||
# Abort Socket Mode WSS handshake after this many seconds. REST auth_test can still
|
||||
# succeed while WSS blocks (firewall / region). slack-sdk does not apply HTTP(S)_PROXY
|
||||
# to websockets.connect — see slack_sdk.socket_mode.websockets.SocketModeClient.connect.
|
||||
SLACK_SOCKET_CONNECT_TIMEOUT_S = 45.0
|
||||
_HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
|
||||
|
||||
|
||||
class SlackChannel(BaseChannel):
|
||||
"""Slack channel using Socket Mode."""
|
||||
|
||||
name = "slack"
|
||||
display_name = "Slack"
|
||||
_SLACK_ID_RE = re.compile(r"^[CDGUW][A-Z0-9]{2,}$")
|
||||
_SLACK_CHANNEL_REF_RE = re.compile(r"^<#([A-Z0-9]+)(?:\|[^>]+)?>$")
|
||||
_SLACK_USER_REF_RE = re.compile(r"^<@([A-Z0-9]+)(?:\|[^>]+)?>$")
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return SlackConfig().model_dump(by_alias=True)
|
||||
|
||||
_THREAD_CONTEXT_CACHE_LIMIT = 10_000
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = SlackConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: SlackConfig = config
|
||||
self._web_client: AsyncWebClient | None = None
|
||||
self._socket_client: SocketModeClient | None = None
|
||||
self._bot_user_id: str | None = None
|
||||
self._target_cache: dict[str, str] = {}
|
||||
self._thread_context_attempted: set[str] = set()
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the Slack Socket Mode client."""
|
||||
if not self.config.bot_token or not self.config.app_token:
|
||||
self.logger.error("bot/app token not configured")
|
||||
return
|
||||
if self.config.mode != "socket":
|
||||
self.logger.error("Unsupported mode: {}", self.config.mode)
|
||||
return
|
||||
|
||||
self._running = True
|
||||
|
||||
self._web_client = AsyncWebClient(token=self.config.bot_token)
|
||||
self._socket_client = SocketModeClient(
|
||||
app_token=self.config.app_token,
|
||||
web_client=self._web_client,
|
||||
)
|
||||
|
||||
self._socket_client.socket_mode_request_listeners.append(self._on_socket_request)
|
||||
|
||||
# Resolve bot user ID for mention handling
|
||||
try:
|
||||
auth = await self._web_client.auth_test()
|
||||
self._bot_user_id = auth.get("user_id")
|
||||
self.logger.info("bot connected as {}", self._bot_user_id)
|
||||
except Exception as e:
|
||||
self.logger.warning("auth_test failed: {}", e)
|
||||
|
||||
self.logger.info("Starting Socket Mode client...")
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._socket_client.connect(),
|
||||
timeout=SLACK_SOCKET_CONNECT_TIMEOUT_S,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
self.logger.error(
|
||||
"Slack Socket Mode WebSocket handshake timed out after {:.0f}s. "
|
||||
"auth_test uses HTTPS and may still succeed while WSS is blocked. "
|
||||
"Check outbound access to Slack WebSockets; slack-sdk Socket Mode "
|
||||
"does not apply HTTP(S)_PROXY to websockets.connect.",
|
||||
SLACK_SOCKET_CONNECT_TIMEOUT_S,
|
||||
)
|
||||
await self.stop()
|
||||
raise RuntimeError("Slack Socket Mode WebSocket connect timed out") from None
|
||||
|
||||
self.logger.info("Slack Socket Mode WebSocket connected (events enabled)")
|
||||
|
||||
while self._running:
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the Slack client."""
|
||||
self._running = False
|
||||
if self._socket_client:
|
||||
try:
|
||||
await self._socket_client.close()
|
||||
except Exception as e:
|
||||
self.logger.warning("socket close failed: {}", e)
|
||||
self._socket_client = None
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a message through Slack."""
|
||||
if not self._web_client:
|
||||
self.logger.warning("client not running")
|
||||
return
|
||||
try:
|
||||
target_chat_id = await self._resolve_target_chat_id(msg.chat_id)
|
||||
slack_meta = msg.metadata.get("slack", {}) if msg.metadata else {}
|
||||
thread_ts = slack_meta.get("thread_ts")
|
||||
origin_chat_id = str((slack_meta.get("event", {}) or {}).get("channel") or msg.chat_id)
|
||||
# Reply in the same thread the inbound message belongs to (works
|
||||
# for both real channel threads and DM threads). When the agent
|
||||
# is forwarding to a different channel, drop thread_ts because it
|
||||
# only makes sense within the originating conversation.
|
||||
thread_ts_param = thread_ts if thread_ts and target_chat_id == origin_chat_id else None
|
||||
|
||||
is_progress = (msg.metadata or {}).get("_progress", False)
|
||||
if is_progress and not msg.content:
|
||||
pass # skip empty progress messages (e.g. tool-event-only updates)
|
||||
elif msg.content or not (msg.media or []):
|
||||
mrkdwn = self._to_mrkdwn(msg.content) if msg.content else " "
|
||||
buttons = getattr(msg, "buttons", None) or []
|
||||
chunks = split_message(mrkdwn, SLACK_MAX_MESSAGE_LEN)
|
||||
for index, chunk in enumerate(chunks):
|
||||
kwargs: dict[str, Any] = dict(
|
||||
channel=target_chat_id, text=chunk, thread_ts=thread_ts_param,
|
||||
)
|
||||
if buttons and index == len(chunks) - 1:
|
||||
kwargs["blocks"] = self._build_button_blocks(chunk, buttons)
|
||||
await self._web_client.chat_postMessage(**kwargs)
|
||||
|
||||
for media_path in msg.media or []:
|
||||
try:
|
||||
await self._web_client.files_upload_v2(
|
||||
channel=target_chat_id,
|
||||
file=media_path,
|
||||
thread_ts=thread_ts_param,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to upload file {}", media_path)
|
||||
|
||||
# Update reaction emoji when the final (non-progress) response is sent
|
||||
if not (msg.metadata or {}).get("_progress"):
|
||||
event = slack_meta.get("event", {})
|
||||
await self._update_react_emoji(origin_chat_id, event.get("ts"))
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Error sending message")
|
||||
raise
|
||||
|
||||
async def _resolve_target_chat_id(self, target: str) -> str:
|
||||
"""Resolve human-friendly Slack targets to concrete IDs when needed."""
|
||||
if not self._web_client:
|
||||
return target
|
||||
|
||||
target = target.strip()
|
||||
if not target:
|
||||
return target
|
||||
|
||||
if match := self._SLACK_CHANNEL_REF_RE.fullmatch(target):
|
||||
return match.group(1)
|
||||
if match := self._SLACK_USER_REF_RE.fullmatch(target):
|
||||
return await self._open_dm_for_user(match.group(1))
|
||||
if self._SLACK_ID_RE.fullmatch(target):
|
||||
if target.startswith(("U", "W")):
|
||||
return await self._open_dm_for_user(target)
|
||||
return target
|
||||
|
||||
if target.startswith("#"):
|
||||
return await self._resolve_channel_name(target[1:])
|
||||
if target.startswith("@"):
|
||||
return await self._resolve_user_handle(target[1:])
|
||||
|
||||
try:
|
||||
return await self._resolve_channel_name(target)
|
||||
except ValueError:
|
||||
return await self._resolve_user_handle(target)
|
||||
|
||||
async def _resolve_channel_name(self, name: str) -> str:
|
||||
normalized = self._normalize_target_name(name)
|
||||
if not normalized:
|
||||
raise ValueError("Slack target channel name is empty")
|
||||
|
||||
cache_key = f"channel:{normalized}"
|
||||
if cache_key in self._target_cache:
|
||||
return self._target_cache[cache_key]
|
||||
|
||||
cursor: str | None = None
|
||||
while True:
|
||||
response = await self._web_client.conversations_list(
|
||||
types="public_channel,private_channel",
|
||||
exclude_archived=True,
|
||||
limit=200,
|
||||
cursor=cursor,
|
||||
)
|
||||
for channel in response.get("channels", []):
|
||||
if self._normalize_target_name(str(channel.get("name") or "")) == normalized:
|
||||
channel_id = str(channel.get("id") or "")
|
||||
if channel_id:
|
||||
self._target_cache[cache_key] = channel_id
|
||||
return channel_id
|
||||
cursor = ((response.get("response_metadata") or {}).get("next_cursor") or "").strip()
|
||||
if not cursor:
|
||||
break
|
||||
|
||||
raise ValueError(
|
||||
f"Slack channel '{name}' was not found. Use a joined channel name like "
|
||||
f"'#general' or a concrete channel ID."
|
||||
)
|
||||
|
||||
async def _resolve_user_handle(self, handle: str) -> str:
|
||||
normalized = self._normalize_target_name(handle)
|
||||
if not normalized:
|
||||
raise ValueError("Slack target user handle is empty")
|
||||
|
||||
cache_key = f"user:{normalized}"
|
||||
if cache_key in self._target_cache:
|
||||
return self._target_cache[cache_key]
|
||||
|
||||
cursor: str | None = None
|
||||
while True:
|
||||
response = await self._web_client.users_list(limit=200, cursor=cursor)
|
||||
for member in response.get("members", []):
|
||||
if self._member_matches_handle(member, normalized):
|
||||
user_id = str(member.get("id") or "")
|
||||
if not user_id:
|
||||
continue
|
||||
dm_id = await self._open_dm_for_user(user_id)
|
||||
self._target_cache[cache_key] = dm_id
|
||||
return dm_id
|
||||
cursor = ((response.get("response_metadata") or {}).get("next_cursor") or "").strip()
|
||||
if not cursor:
|
||||
break
|
||||
|
||||
raise ValueError(
|
||||
f"Slack user '{handle}' was not found. Use '@name' or a concrete DM/channel ID."
|
||||
)
|
||||
|
||||
async def _open_dm_for_user(self, user_id: str) -> str:
|
||||
response = await self._web_client.conversations_open(users=user_id)
|
||||
channel_id = str(((response.get("channel") or {}).get("id")) or "")
|
||||
if not channel_id:
|
||||
raise ValueError(f"Slack DM target for user '{user_id}' could not be opened.")
|
||||
return channel_id
|
||||
|
||||
@staticmethod
|
||||
def _normalize_target_name(value: str) -> str:
|
||||
return value.strip().lstrip("#@").lower()
|
||||
|
||||
@classmethod
|
||||
def _member_matches_handle(cls, member: dict[str, Any], normalized: str) -> bool:
|
||||
profile = member.get("profile") or {}
|
||||
candidates = {
|
||||
str(member.get("name") or ""),
|
||||
str(profile.get("display_name") or ""),
|
||||
str(profile.get("display_name_normalized") or ""),
|
||||
str(profile.get("real_name") or ""),
|
||||
str(profile.get("real_name_normalized") or ""),
|
||||
}
|
||||
return normalized in {cls._normalize_target_name(candidate) for candidate in candidates if candidate}
|
||||
|
||||
async def _on_socket_request(
|
||||
self,
|
||||
client: SocketModeClient,
|
||||
req: SocketModeRequest,
|
||||
) -> None:
|
||||
"""Handle incoming Socket Mode requests."""
|
||||
if req.type == "interactive":
|
||||
await self._on_block_action(client, req)
|
||||
return
|
||||
if req.type != "events_api":
|
||||
return
|
||||
|
||||
# Acknowledge right away
|
||||
await client.send_socket_mode_response(
|
||||
SocketModeResponse(envelope_id=req.envelope_id)
|
||||
)
|
||||
|
||||
payload = req.payload or {}
|
||||
event = payload.get("event") or {}
|
||||
event_type = event.get("type")
|
||||
|
||||
# Handle app mentions or plain messages
|
||||
if event_type not in ("message", "app_mention"):
|
||||
return
|
||||
|
||||
sender_id = event.get("user")
|
||||
chat_id = event.get("channel")
|
||||
|
||||
subtype = event.get("subtype")
|
||||
# Slack uses subtype=file_share for user messages with attachments.
|
||||
# Ignore other subtypes such as bot_message / message_changed / deleted.
|
||||
if subtype and subtype != "file_share":
|
||||
return
|
||||
if self._bot_user_id and sender_id == self._bot_user_id:
|
||||
return
|
||||
|
||||
# Avoid double-processing: Slack sends both `message` and `app_mention`
|
||||
# for mentions in channels. Prefer `app_mention`.
|
||||
text = event.get("text") or ""
|
||||
if event_type == "message" and self._bot_user_id and f"<@{self._bot_user_id}>" in text:
|
||||
return
|
||||
|
||||
# Debug: log basic event shape
|
||||
self.logger.debug(
|
||||
"event: type={} subtype={} user={} channel={} channel_type={} text={}",
|
||||
event_type,
|
||||
subtype,
|
||||
sender_id,
|
||||
chat_id,
|
||||
event.get("channel_type"),
|
||||
text[:80],
|
||||
)
|
||||
if not sender_id or not chat_id:
|
||||
return
|
||||
|
||||
channel_type = event.get("channel_type") or ""
|
||||
|
||||
if not self._is_allowed(sender_id, chat_id, channel_type):
|
||||
if channel_type == "im" and self.config.dm.enabled:
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
content="",
|
||||
is_dm=True,
|
||||
)
|
||||
return
|
||||
|
||||
if channel_type != "im" and not self._should_respond_in_channel(event_type, text, chat_id):
|
||||
return
|
||||
|
||||
text = self._strip_bot_mention(text)
|
||||
|
||||
event_ts = event.get("ts")
|
||||
raw_thread_ts = event.get("thread_ts")
|
||||
thread_ts = raw_thread_ts
|
||||
# In DMs we don't auto-open a thread on top-level messages (it would
|
||||
# bury replies under "1 reply"). But if the user explicitly opened a
|
||||
# thread inside the DM, raw_thread_ts is set and we honor it.
|
||||
if (
|
||||
self.config.reply_in_thread
|
||||
and not thread_ts
|
||||
and channel_type != "im"
|
||||
):
|
||||
thread_ts = event_ts
|
||||
# Add :eyes: reaction to the triggering message (best-effort)
|
||||
try:
|
||||
if self._web_client and event.get("ts"):
|
||||
await self._web_client.reactions_add(
|
||||
channel=chat_id,
|
||||
name=self.config.react_emoji,
|
||||
timestamp=event.get("ts"),
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.debug("reactions_add failed: {}", e)
|
||||
|
||||
# Thread-scoped session key whenever the user is in a real thread
|
||||
# (raw_thread_ts is set). DM threads get their own session, separate
|
||||
# from the DM root, so context doesn't bleed across thread boundaries.
|
||||
session_key = (
|
||||
f"slack:{chat_id}:{thread_ts}" if thread_ts and raw_thread_ts else None
|
||||
)
|
||||
media_paths: list[str] = []
|
||||
file_markers: list[str] = []
|
||||
for file_info in event.get("files") or []:
|
||||
if not isinstance(file_info, dict):
|
||||
continue
|
||||
file_path, marker = await self._download_slack_file(file_info)
|
||||
if file_path:
|
||||
media_paths.append(file_path)
|
||||
if marker:
|
||||
file_markers.append(marker)
|
||||
|
||||
is_slash = text.strip().startswith("/")
|
||||
content = text if is_slash else await self._with_thread_context(
|
||||
text,
|
||||
chat_id=chat_id,
|
||||
channel_type=channel_type,
|
||||
thread_ts=thread_ts,
|
||||
raw_thread_ts=raw_thread_ts,
|
||||
current_ts=event_ts,
|
||||
)
|
||||
if file_markers:
|
||||
content = "\n".join(part for part in [content, *file_markers] if part)
|
||||
if not content and not media_paths:
|
||||
return
|
||||
|
||||
try:
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
content=content,
|
||||
media=media_paths,
|
||||
metadata={
|
||||
"slack": {
|
||||
"event": event,
|
||||
"thread_ts": thread_ts,
|
||||
"channel_type": channel_type,
|
||||
},
|
||||
},
|
||||
session_key=session_key,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Error handling message from {}", sender_id)
|
||||
|
||||
async def _download_slack_file(self, file_info: dict[str, Any]) -> tuple[str | None, str]:
|
||||
"""Download a Slack private file to the local media directory."""
|
||||
file_id = str(file_info.get("id") or "file")
|
||||
name = str(
|
||||
file_info.get("name")
|
||||
or file_info.get("title")
|
||||
or file_info.get("id")
|
||||
or "slack-file"
|
||||
)
|
||||
marker_type = "image" if str(file_info.get("mimetype") or "").startswith("image/") else "file"
|
||||
marker = f"[{marker_type}: {name}]"
|
||||
url = str(file_info.get("url_private_download") or file_info.get("url_private") or "")
|
||||
if not url:
|
||||
return None, self._download_failure_marker(marker_type, name, "missing download url")
|
||||
if not self.config.bot_token:
|
||||
return None, self._download_failure_marker(marker_type, name, "missing bot token")
|
||||
|
||||
filename = safe_filename(f"{file_id}_{name}")
|
||||
path = Path(get_media_dir("slack")) / filename
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=SLACK_DOWNLOAD_TIMEOUT, follow_redirects=True) as client:
|
||||
response = await client.get(
|
||||
url,
|
||||
headers={"Authorization": f"Bearer {self.config.bot_token}"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
if self._looks_like_html_download(response):
|
||||
raise ValueError("Slack returned HTML instead of file content")
|
||||
path.write_bytes(response.content)
|
||||
return str(path), marker
|
||||
except Exception as e:
|
||||
self.logger.warning("Failed to download file {}: {}", file_id, e)
|
||||
return None, self._download_failure_marker(marker_type, name, "download failed")
|
||||
|
||||
@staticmethod
|
||||
def _download_failure_marker(marker_type: str, name: str, reason: str) -> str:
|
||||
return (
|
||||
f"[{marker_type}: {name}: {reason}; not available to vibe-trading. "
|
||||
"Check Slack files:read scope, reinstall the Slack app, and ensure the bot can access the file.]"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _looks_like_html_download(response: httpx.Response) -> bool:
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
if "text/html" in content_type:
|
||||
return True
|
||||
preview = response.content[:256].lstrip().lower()
|
||||
return preview.startswith(_HTML_DOWNLOAD_PREFIXES)
|
||||
|
||||
async def _on_block_action(self, client: SocketModeClient, req: SocketModeRequest) -> None:
|
||||
"""Handle button clicks from inline action buttons."""
|
||||
await client.send_socket_mode_response(SocketModeResponse(envelope_id=req.envelope_id))
|
||||
payload = req.payload or {}
|
||||
actions = payload.get("actions") or []
|
||||
if not actions:
|
||||
return
|
||||
value = str(actions[0].get("value") or "")
|
||||
user_info = payload.get("user") or {}
|
||||
sender_id = str(user_info.get("id") or "")
|
||||
channel_info = payload.get("channel") or {}
|
||||
chat_id = str(channel_info.get("id") or "")
|
||||
if not sender_id or not chat_id or not value:
|
||||
return
|
||||
message_info = payload.get("message") or {}
|
||||
thread_ts = message_info.get("thread_ts") or message_info.get("ts")
|
||||
channel_type = self._infer_channel_type(chat_id)
|
||||
if not self._is_allowed(sender_id, chat_id, channel_type):
|
||||
return
|
||||
session_key = f"slack:{chat_id}:{thread_ts}" if thread_ts else None
|
||||
try:
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
content=value,
|
||||
metadata={"slack": {"thread_ts": thread_ts, "channel_type": channel_type}},
|
||||
session_key=session_key,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Error handling button click from {}", sender_id)
|
||||
|
||||
async def _with_thread_context(
|
||||
self,
|
||||
text: str,
|
||||
*,
|
||||
chat_id: str,
|
||||
channel_type: str,
|
||||
thread_ts: str | None,
|
||||
raw_thread_ts: str | None,
|
||||
current_ts: str | None,
|
||||
) -> str:
|
||||
"""Include thread history the first time the bot is pulled into a Slack thread."""
|
||||
del channel_type # DM and channel threads are both fetched via conversations.replies
|
||||
if (
|
||||
not self.config.include_thread_context
|
||||
or not self._web_client
|
||||
or not raw_thread_ts
|
||||
or not thread_ts
|
||||
or current_ts == thread_ts
|
||||
):
|
||||
return text
|
||||
|
||||
key = f"{chat_id}:{thread_ts}"
|
||||
if key in self._thread_context_attempted:
|
||||
return text
|
||||
if len(self._thread_context_attempted) >= self._THREAD_CONTEXT_CACHE_LIMIT:
|
||||
self._thread_context_attempted.clear()
|
||||
self._thread_context_attempted.add(key)
|
||||
|
||||
try:
|
||||
response = await self._web_client.conversations_replies(
|
||||
channel=chat_id,
|
||||
ts=thread_ts,
|
||||
limit=max(1, self.config.thread_context_limit),
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning("thread context unavailable for {}: {}", key, e)
|
||||
return text
|
||||
|
||||
lines = self._format_thread_context(
|
||||
response.get("messages", []),
|
||||
current_ts=current_ts,
|
||||
)
|
||||
if not lines:
|
||||
return text
|
||||
return "Slack thread context before this mention:\n" + "\n".join(lines) + f"\n\nCurrent message:\n{text}"
|
||||
|
||||
def _format_thread_context(self, messages: list[dict[str, Any]], *, current_ts: str | None) -> list[str]:
|
||||
lines: list[str] = []
|
||||
for item in messages:
|
||||
if item.get("ts") == current_ts:
|
||||
continue
|
||||
if item.get("subtype"):
|
||||
continue
|
||||
sender = str(item.get("user") or item.get("bot_id") or "unknown")
|
||||
is_bot = self._bot_user_id is not None and sender == self._bot_user_id
|
||||
label = "bot" if is_bot else f"<@{sender}>"
|
||||
text = str(item.get("text") or "").strip()
|
||||
if not text:
|
||||
continue
|
||||
text = self._strip_bot_mention(text)
|
||||
if len(text) > 500:
|
||||
text = text[:500] + "…"
|
||||
lines.append(f"- {label}: {text}")
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def _build_button_blocks(text: str, buttons: list[list[str]]) -> list[dict[str, Any]]:
|
||||
"""Build Slack Block Kit blocks with action buttons."""
|
||||
blocks: list[dict[str, Any]] = [
|
||||
{"type": "section", "text": {"type": "mrkdwn", "text": text[:3000]}},
|
||||
]
|
||||
elements = []
|
||||
for row in buttons:
|
||||
for label in row:
|
||||
elements.append({
|
||||
"type": "button",
|
||||
"text": {"type": "plain_text", "text": label[:75]},
|
||||
"value": label[:75],
|
||||
"action_id": f"btn_{label[:50]}",
|
||||
})
|
||||
if elements:
|
||||
blocks.append({"type": "actions", "elements": elements[:25]})
|
||||
return blocks
|
||||
|
||||
async def _update_react_emoji(self, chat_id: str, ts: str | None) -> None:
|
||||
"""Remove the in-progress reaction and optionally add a done reaction."""
|
||||
if not self._web_client or not ts:
|
||||
return
|
||||
try:
|
||||
await self._web_client.reactions_remove(
|
||||
channel=chat_id,
|
||||
name=self.config.react_emoji,
|
||||
timestamp=ts,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.debug("reactions_remove failed: {}", e)
|
||||
if self.config.done_emoji:
|
||||
try:
|
||||
await self._web_client.reactions_add(
|
||||
channel=chat_id,
|
||||
name=self.config.done_emoji,
|
||||
timestamp=ts,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.debug("done reaction failed: {}", e)
|
||||
|
||||
def _is_allowed(self, sender_id: str, chat_id: str, channel_type: str) -> bool:
|
||||
if channel_type == "im":
|
||||
if not self.config.dm.enabled:
|
||||
return False
|
||||
if self.config.dm.policy == "allowlist":
|
||||
return sender_id in self.config.dm.allow_from or is_approved(self.name, sender_id)
|
||||
return True
|
||||
|
||||
# Group / channel messages
|
||||
if self.config.group_policy == "allowlist":
|
||||
return chat_id in self.config.group_allow_from
|
||||
return True
|
||||
|
||||
def _is_mention(self, event_type: str, text: str) -> bool:
|
||||
if event_type == "app_mention":
|
||||
return True
|
||||
return self._bot_user_id is not None and f"<@{self._bot_user_id}>" in text
|
||||
|
||||
def _should_respond_in_channel(self, event_type: str, text: str, chat_id: str) -> bool:
|
||||
if self.config.group_policy == "open":
|
||||
return True
|
||||
if self.config.group_policy == "mention":
|
||||
return self._is_mention(event_type, text)
|
||||
if self.config.group_policy == "allowlist":
|
||||
if chat_id not in self.config.group_allow_from:
|
||||
return False
|
||||
if self.config.group_require_mention:
|
||||
return self._is_mention(event_type, text)
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_allowed(self, sender_id: str) -> bool:
|
||||
# Slack needs channel-aware policy checks, so _on_socket_request and
|
||||
# _on_block_action call _is_allowed before handing off to BaseChannel.
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _infer_channel_type(chat_id: str) -> str:
|
||||
if chat_id.startswith("D"):
|
||||
return "im"
|
||||
if chat_id.startswith("G"):
|
||||
return "group"
|
||||
return "channel"
|
||||
|
||||
def _strip_bot_mention(self, text: str) -> str:
|
||||
if not text or not self._bot_user_id:
|
||||
return text
|
||||
return re.sub(rf"<@{re.escape(self._bot_user_id)}>\s*", "", text).strip()
|
||||
|
||||
_TABLE_RE = re.compile(r"(?m)^\|.*\|$(?:\n\|[\s:|-]*\|$)(?:\n\|.*\|$)*")
|
||||
_CODE_FENCE_RE = re.compile(r"```[\s\S]*?```")
|
||||
_INLINE_CODE_RE = re.compile(r"`[^`]+`")
|
||||
_LEFTOVER_BOLD_RE = re.compile(r"\*\*(.+?)\*\*")
|
||||
_LEFTOVER_HEADER_RE = re.compile(r"^#{1,6}\s+(.+)$", re.MULTILINE)
|
||||
_BARE_URL_RE = re.compile(r"(?<![|<])(https?://\S+)")
|
||||
|
||||
@classmethod
|
||||
def _to_mrkdwn(cls, text: str) -> str:
|
||||
"""Convert Markdown to Slack mrkdwn, including tables."""
|
||||
if not text:
|
||||
return ""
|
||||
text = cls._TABLE_RE.sub(cls._convert_table, text)
|
||||
return cls._fixup_mrkdwn(slackify_markdown(text)).rstrip("\n")
|
||||
|
||||
@classmethod
|
||||
def _fixup_mrkdwn(cls, text: str) -> str:
|
||||
"""Fix markdown artifacts that slackify_markdown misses."""
|
||||
code_blocks: list[str] = []
|
||||
|
||||
def _save_code(m: re.Match) -> str:
|
||||
code_blocks.append(m.group(0))
|
||||
return f"\x00CB{len(code_blocks) - 1}\x00"
|
||||
|
||||
text = cls._CODE_FENCE_RE.sub(_save_code, text)
|
||||
text = cls._INLINE_CODE_RE.sub(_save_code, text)
|
||||
text = cls._LEFTOVER_BOLD_RE.sub(r"*\1*", text)
|
||||
text = cls._LEFTOVER_HEADER_RE.sub(r"*\1*", text)
|
||||
text = cls._BARE_URL_RE.sub(lambda m: m.group(0).replace("&", "&"), text)
|
||||
|
||||
for i, block in enumerate(code_blocks):
|
||||
text = text.replace(f"\x00CB{i}\x00", block)
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _convert_table(match: re.Match) -> str:
|
||||
"""Convert a Markdown table to a Slack-readable list."""
|
||||
lines = [ln.strip() for ln in match.group(0).strip().splitlines() if ln.strip()]
|
||||
if len(lines) < 2:
|
||||
return match.group(0)
|
||||
headers = [h.strip() for h in lines[0].strip("|").split("|")]
|
||||
start = 2 if re.fullmatch(r"[|\s:\-]+", lines[1]) else 1
|
||||
rows: list[str] = []
|
||||
for line in lines[start:]:
|
||||
cells = [c.strip() for c in line.strip("|").split("|")]
|
||||
cells = (cells + [""] * len(headers))[: len(headers)]
|
||||
parts = [f"**{headers[i]}**: {cells[i]}" for i in range(len(headers)) if cells[i]]
|
||||
if parts:
|
||||
rows.append(" · ".join(parts))
|
||||
return "\n".join(rows)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,156 @@
|
||||
"""Utility helpers for channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import re
|
||||
import socket
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from src.config.paths import get_data_dir
|
||||
|
||||
_UNSAFE_CHARS = re.compile(r"[/\\:*?\"<>|]")
|
||||
|
||||
|
||||
def get_media_dir(channel_name: str) -> Path:
|
||||
"""Return a media directory for *channel_name* under the VT data dir."""
|
||||
p = get_data_dir() / channel_name
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
|
||||
|
||||
def get_runtime_subdir(name: str) -> Path:
|
||||
"""Return a runtime subdirectory for *name* under the VT data dir."""
|
||||
p = get_data_dir() / "runtime" / name
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
|
||||
|
||||
def validate_resolved_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
"""Validate a URL then resolve and check its IP addresses.
|
||||
|
||||
Thin wrapper around :func:`validate_url_target` for callers that want
|
||||
the same behavior under a legacy name.
|
||||
"""
|
||||
return validate_url_target(url, allow_loopback=allow_loopback)
|
||||
|
||||
|
||||
def is_path_within(path: str | Path, root: str | Path) -> bool:
|
||||
"""Check whether *path* resides within *root* directory."""
|
||||
try:
|
||||
Path(path).resolve().relative_to(Path(root).resolve())
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def split_message(content: str, max_len: int = 2000) -> list[str]:
|
||||
"""Split content into chunks within max_len, preferring line breaks.
|
||||
|
||||
Args:
|
||||
content: The text content to split.
|
||||
max_len: Maximum length per chunk (default 2000 for Discord compatibility).
|
||||
|
||||
Returns:
|
||||
List of message chunks, each within max_len.
|
||||
"""
|
||||
if not content:
|
||||
return []
|
||||
if len(content) <= max_len:
|
||||
return [content]
|
||||
chunks: list[str] = []
|
||||
while content:
|
||||
if len(content) <= max_len:
|
||||
chunks.append(content)
|
||||
break
|
||||
cut = content[:max_len]
|
||||
# Try to break at newline first, then space, then hard break
|
||||
pos = cut.rfind("\n")
|
||||
if pos <= 0:
|
||||
pos = cut.rfind(" ")
|
||||
if pos <= 0:
|
||||
pos = max_len
|
||||
chunks.append(content[:pos])
|
||||
content = content[pos:].lstrip()
|
||||
return chunks
|
||||
|
||||
|
||||
def safe_filename(name: str) -> str:
|
||||
"""Replace unsafe path characters with underscores."""
|
||||
return _UNSAFE_CHARS.sub("_", name).strip()
|
||||
|
||||
|
||||
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
"""Validate a URL is safe to fetch: scheme, hostname, and resolved IPs.
|
||||
|
||||
``allow_loopback`` is intentionally narrow: it only permits literal
|
||||
loopback hosts (localhost, 127.0.0.0/8, ::1) when every resolved address is
|
||||
loopback. It does not allow RFC1918, link-local, metadata, or public DNS
|
||||
names that happen to resolve to loopback.
|
||||
|
||||
Returns (ok, error_message). When ok is True, error_message is empty.
|
||||
"""
|
||||
try:
|
||||
p = urlparse(url)
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
if p.scheme not in ("http", "https"):
|
||||
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'"
|
||||
if not p.netloc:
|
||||
return False, "Missing domain"
|
||||
|
||||
hostname = p.hostname
|
||||
if not hostname:
|
||||
return False, "Missing hostname"
|
||||
|
||||
try:
|
||||
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
||||
except socket.gaierror:
|
||||
return False, f"Cannot resolve hostname: {hostname}"
|
||||
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||
for info in infos:
|
||||
try:
|
||||
addr = ipaddress.ip_address(info[4][0])
|
||||
except ValueError:
|
||||
continue
|
||||
addrs.append(addr)
|
||||
|
||||
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
|
||||
return True, ""
|
||||
|
||||
for addr in addrs:
|
||||
if _is_private(addr):
|
||||
return False, f"Blocked: {hostname} resolves to private/internal address {addr}"
|
||||
|
||||
return True, ""
|
||||
|
||||
|
||||
def _is_allowed_loopback_target(
|
||||
hostname: str,
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address],
|
||||
) -> bool:
|
||||
"""Check that *hostname* is a literal loopback host and every addr is loopback."""
|
||||
if hostname.lower() not in ("localhost", "localhost."):
|
||||
# Check for literal "127.x.y.z" or "[::1]"
|
||||
is_loopback_literal = False
|
||||
for addr in addrs:
|
||||
if addr.is_loopback:
|
||||
is_loopback_literal = True
|
||||
else:
|
||||
return False
|
||||
return is_loopback_literal
|
||||
return all(addr.is_loopback for addr in addrs) if addrs else False
|
||||
|
||||
|
||||
def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||||
"""Check whether *addr* is private, link-local, or loopback."""
|
||||
if addr.is_loopback or addr.is_link_local or addr.is_private:
|
||||
return True
|
||||
# Also block well-known metadata endpoints
|
||||
if isinstance(addr, ipaddress.IPv4Address):
|
||||
return False
|
||||
# IPv6: block unspecified, multicast, unique-local
|
||||
return addr.is_multicast or addr.is_unspecified
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,554 @@
|
||||
"""WeCom (Enterprise WeChat) channel implementation using wecom_aibot_sdk."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import os
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_media_dir
|
||||
from pydantic import BaseModel
|
||||
|
||||
WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None
|
||||
|
||||
# Upload safety limits (matching QQ channel defaults)
|
||||
WECOM_UPLOAD_MAX_BYTES = 1024 * 1024 * 200 # 200MB
|
||||
|
||||
# Replace unsafe characters with "_", keep Chinese and common safe punctuation.
|
||||
_SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]()【】\u4e00-\u9fff]+", re.UNICODE)
|
||||
|
||||
|
||||
def _sanitize_filename(name: str) -> str:
|
||||
"""Sanitize filename to avoid traversal and problematic chars."""
|
||||
name = (name or "").strip()
|
||||
name = Path(name).name
|
||||
name = _SAFE_NAME_RE.sub("_", name).strip("._ ")
|
||||
return name
|
||||
|
||||
|
||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
|
||||
_VIDEO_EXTS = {".mp4", ".avi", ".mov"}
|
||||
_AUDIO_EXTS = {".amr", ".mp3", ".wav", ".ogg"}
|
||||
|
||||
|
||||
def _guess_wecom_media_type(filename: str) -> str:
|
||||
"""Classify file extension as WeCom media_type string."""
|
||||
ext = Path(filename).suffix.lower()
|
||||
if ext in _IMAGE_EXTS:
|
||||
return "image"
|
||||
if ext in _VIDEO_EXTS:
|
||||
return "video"
|
||||
if ext in _AUDIO_EXTS:
|
||||
return "voice"
|
||||
return "file"
|
||||
|
||||
class WecomConfig(BaseModel):
|
||||
"""WeCom (Enterprise WeChat) AI Bot channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
bot_id: str = ""
|
||||
secret: str = ""
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
welcome_message: str = ""
|
||||
|
||||
|
||||
# Message type display mapping
|
||||
MSG_TYPE_MAP = {
|
||||
"image": "[image]",
|
||||
"voice": "[voice]",
|
||||
"file": "[file]",
|
||||
"mixed": "[mixed content]",
|
||||
}
|
||||
|
||||
|
||||
class WecomChannel(BaseChannel):
|
||||
"""
|
||||
WeCom (Enterprise WeChat) channel using WebSocket long connection.
|
||||
|
||||
Uses WebSocket to receive events - no public IP or webhook required.
|
||||
|
||||
Requires:
|
||||
- Bot ID and Secret from WeCom AI Bot platform
|
||||
"""
|
||||
|
||||
name = "wecom"
|
||||
display_name = "WeCom"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return WecomConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
if isinstance(config, dict):
|
||||
config = WecomConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
self.config: WecomConfig = config
|
||||
self._client: Any = None
|
||||
self._processed_message_ids: OrderedDict[str, None] = OrderedDict()
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._generate_req_id = None
|
||||
# Store frame headers for each chat to enable replies
|
||||
self._chat_frames: dict[str, Any] = {}
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the WeCom bot with WebSocket long connection."""
|
||||
if not WECOM_AVAILABLE:
|
||||
self.logger.error("SDK not installed. Run: pip install vibe-trading-ai[wecom]")
|
||||
return
|
||||
|
||||
if not self.config.bot_id or not self.config.secret:
|
||||
self.logger.error("bot_id and secret not configured")
|
||||
return
|
||||
|
||||
from wecom_aibot_sdk import WSClient, generate_req_id
|
||||
|
||||
self._running = True
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._generate_req_id = generate_req_id
|
||||
|
||||
# Create WebSocket client
|
||||
self._client = WSClient({
|
||||
"bot_id": self.config.bot_id,
|
||||
"secret": self.config.secret,
|
||||
"reconnect_interval": 1000,
|
||||
"max_reconnect_attempts": -1, # Infinite reconnect
|
||||
"heartbeat_interval": 30000,
|
||||
})
|
||||
|
||||
# Register event handlers
|
||||
self._client.on("connected", self._on_connected)
|
||||
self._client.on("authenticated", self._on_authenticated)
|
||||
self._client.on("disconnected", self._on_disconnected)
|
||||
self._client.on("error", self._on_error)
|
||||
self._client.on("message.text", self._on_text_message)
|
||||
self._client.on("message.image", self._on_image_message)
|
||||
self._client.on("message.voice", self._on_voice_message)
|
||||
self._client.on("message.file", self._on_file_message)
|
||||
self._client.on("message.mixed", self._on_mixed_message)
|
||||
self._client.on("event.enter_chat", self._on_enter_chat)
|
||||
|
||||
self.logger.info("bot starting with WebSocket long connection")
|
||||
self.logger.info("No public IP required - using WebSocket to receive events")
|
||||
|
||||
# Connect
|
||||
await self._client.connect_async()
|
||||
|
||||
# Keep running until stopped
|
||||
while self._running:
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the WeCom bot."""
|
||||
self._running = False
|
||||
if self._client:
|
||||
await self._client.disconnect()
|
||||
self.logger.info("bot stopped")
|
||||
|
||||
async def _on_connected(self, frame: Any) -> None:
|
||||
"""Handle WebSocket connected event."""
|
||||
self.logger.info("WebSocket connected")
|
||||
|
||||
async def _on_authenticated(self, frame: Any) -> None:
|
||||
"""Handle authentication success event."""
|
||||
self.logger.info("authenticated successfully")
|
||||
|
||||
async def _on_disconnected(self, frame: Any) -> None:
|
||||
"""Handle WebSocket disconnected event."""
|
||||
reason = frame.body if hasattr(frame, 'body') else str(frame)
|
||||
self.logger.warning("WebSocket disconnected: {}", reason)
|
||||
|
||||
async def _on_error(self, frame: Any) -> None:
|
||||
"""Handle error event."""
|
||||
self.logger.error("error: {}", frame)
|
||||
|
||||
async def _on_text_message(self, frame: Any) -> None:
|
||||
"""Handle text message."""
|
||||
await self._process_message(frame, "text")
|
||||
|
||||
async def _on_image_message(self, frame: Any) -> None:
|
||||
"""Handle image message."""
|
||||
await self._process_message(frame, "image")
|
||||
|
||||
async def _on_voice_message(self, frame: Any) -> None:
|
||||
"""Handle voice message."""
|
||||
await self._process_message(frame, "voice")
|
||||
|
||||
async def _on_file_message(self, frame: Any) -> None:
|
||||
"""Handle file message."""
|
||||
await self._process_message(frame, "file")
|
||||
|
||||
async def _on_mixed_message(self, frame: Any) -> None:
|
||||
"""Handle mixed content message."""
|
||||
await self._process_message(frame, "mixed")
|
||||
|
||||
async def _on_enter_chat(self, frame: Any) -> None:
|
||||
"""Handle enter_chat event (user opens chat with bot)."""
|
||||
try:
|
||||
# Extract body from WsFrame dataclass or dict
|
||||
if hasattr(frame, 'body'):
|
||||
body = frame.body or {}
|
||||
elif isinstance(frame, dict):
|
||||
body = frame.get("body", frame)
|
||||
else:
|
||||
body = {}
|
||||
|
||||
chat_id = body.get("chatid", "") if isinstance(body, dict) else ""
|
||||
|
||||
if chat_id and not self.is_allowed(chat_id):
|
||||
return
|
||||
|
||||
if chat_id and self.config.welcome_message:
|
||||
await self._client.reply_welcome(frame, {
|
||||
"msgtype": "text",
|
||||
"text": {"content": self.config.welcome_message},
|
||||
})
|
||||
except Exception:
|
||||
self.logger.exception("Error handling enter_chat")
|
||||
|
||||
async def _process_message(self, frame: Any, msg_type: str) -> None:
|
||||
"""Process incoming message and forward to bus."""
|
||||
try:
|
||||
# Extract body from WsFrame dataclass or dict
|
||||
if hasattr(frame, 'body'):
|
||||
body = frame.body or {}
|
||||
elif isinstance(frame, dict):
|
||||
body = frame.get("body", frame)
|
||||
else:
|
||||
body = {}
|
||||
|
||||
# Ensure body is a dict
|
||||
if not isinstance(body, dict):
|
||||
self.logger.warning("Invalid body type: {}", type(body))
|
||||
return
|
||||
|
||||
# Extract message info
|
||||
msg_id = body.get("msgid", "")
|
||||
if not msg_id:
|
||||
msg_id = f"{body.get('chatid', '')}_{body.get('sendertime', '')}"
|
||||
|
||||
# Extract sender info from "from" field (SDK format)
|
||||
from_info = body.get("from", {})
|
||||
sender_id = from_info.get("userid", "unknown") if isinstance(from_info, dict) else "unknown"
|
||||
if not self.is_allowed(sender_id):
|
||||
return
|
||||
|
||||
# Deduplication check
|
||||
if msg_id in self._processed_message_ids:
|
||||
return
|
||||
self._processed_message_ids[msg_id] = None
|
||||
|
||||
# Trim cache
|
||||
while len(self._processed_message_ids) > 1000:
|
||||
self._processed_message_ids.popitem(last=False)
|
||||
|
||||
# For single chat, chatid is the sender's userid
|
||||
# For group chat, chatid is provided in body
|
||||
chat_type = body.get("chattype", "single")
|
||||
chat_id = body.get("chatid", sender_id)
|
||||
|
||||
content_parts = []
|
||||
media_paths: list[str] = []
|
||||
|
||||
if msg_type == "text":
|
||||
text = body.get("text", {}).get("content", "")
|
||||
if text:
|
||||
content_parts.append(text)
|
||||
|
||||
elif msg_type == "image":
|
||||
image_info = body.get("image", {})
|
||||
file_url = image_info.get("url", "")
|
||||
aes_key = image_info.get("aeskey", "")
|
||||
|
||||
if file_url and aes_key:
|
||||
file_path = await self._download_and_save_media(file_url, aes_key, "image")
|
||||
if file_path:
|
||||
filename = os.path.basename(file_path)
|
||||
content_parts.append(f"[image: {filename}]")
|
||||
media_paths.append(file_path)
|
||||
else:
|
||||
content_parts.append("[image: download failed]")
|
||||
else:
|
||||
content_parts.append("[image: download failed]")
|
||||
|
||||
elif msg_type == "voice":
|
||||
voice_info = body.get("voice", {})
|
||||
# Voice message already contains transcribed content from WeCom
|
||||
voice_content = voice_info.get("content", "")
|
||||
if voice_content:
|
||||
content_parts.append(f"[voice] {voice_content}")
|
||||
else:
|
||||
content_parts.append("[voice]")
|
||||
|
||||
elif msg_type == "file":
|
||||
file_info = body.get("file", {})
|
||||
file_url = file_info.get("url", "")
|
||||
aes_key = file_info.get("aeskey", "")
|
||||
file_name = file_info.get("name") or None
|
||||
|
||||
if file_url and aes_key:
|
||||
file_path = await self._download_and_save_media(file_url, aes_key, "file", file_name)
|
||||
if file_path:
|
||||
display_name = os.path.basename(file_path)
|
||||
content_parts.append(f"[file: {display_name}]")
|
||||
media_paths.append(file_path)
|
||||
else:
|
||||
content_parts.append(f"[file: {file_name or 'unknown'}: download failed]")
|
||||
else:
|
||||
content_parts.append(f"[file: {file_name or 'unknown'}: download failed]")
|
||||
|
||||
elif msg_type == "mixed":
|
||||
# Mixed content contains multiple message items
|
||||
msg_items = body.get("mixed", {}).get("msg_item", [])
|
||||
for item in msg_items:
|
||||
item_type = item.get("msgtype", "")
|
||||
if item_type == "text":
|
||||
text = item.get("text", {}).get("content", "")
|
||||
if text:
|
||||
content_parts.append(text)
|
||||
elif item_type == "image":
|
||||
file_url = item.get("image", {}).get("url", "")
|
||||
aes_key = item.get("image", {}).get("aeskey", "")
|
||||
if file_url and aes_key:
|
||||
file_path = await self._download_and_save_media(file_url, aes_key, "image")
|
||||
if file_path:
|
||||
filename = os.path.basename(file_path)
|
||||
content_parts.append(f"[image: {filename}]")
|
||||
media_paths.append(file_path)
|
||||
else:
|
||||
content_parts.append(MSG_TYPE_MAP.get(item_type, f"[{item_type}]"))
|
||||
|
||||
else:
|
||||
content_parts.append(MSG_TYPE_MAP.get(msg_type, f"[{msg_type}]"))
|
||||
|
||||
content = "\n".join(content_parts) if content_parts else ""
|
||||
|
||||
if not content:
|
||||
return
|
||||
|
||||
# Store frame for this chat to enable replies
|
||||
self._chat_frames[chat_id] = frame
|
||||
|
||||
# Forward to message bus
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
content=content,
|
||||
media=media_paths or None,
|
||||
metadata={
|
||||
"message_id": msg_id,
|
||||
"msg_type": msg_type,
|
||||
"chat_type": chat_type,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Error processing message")
|
||||
|
||||
async def _download_and_save_media(
|
||||
self,
|
||||
file_url: str,
|
||||
aes_key: str,
|
||||
media_type: str,
|
||||
filename: str | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Download and decrypt media from WeCom.
|
||||
|
||||
Returns:
|
||||
file_path or None if download failed
|
||||
"""
|
||||
try:
|
||||
data, fname = await self._client.download_file(file_url, aes_key)
|
||||
|
||||
if not data:
|
||||
self.logger.warning("Failed to download media")
|
||||
return None
|
||||
|
||||
if len(data) > WECOM_UPLOAD_MAX_BYTES:
|
||||
self.logger.warning(
|
||||
"inbound media too large: {} bytes (max {})",
|
||||
len(data),
|
||||
WECOM_UPLOAD_MAX_BYTES,
|
||||
)
|
||||
return None
|
||||
|
||||
media_dir = get_media_dir("wecom")
|
||||
if not filename:
|
||||
filename = fname or f"{media_type}_{hash(file_url) % 100000}"
|
||||
filename = _sanitize_filename(filename)
|
||||
|
||||
file_path = media_dir / filename
|
||||
await asyncio.to_thread(file_path.write_bytes, data)
|
||||
self.logger.debug("Downloaded {} to {}", media_type, file_path)
|
||||
return str(file_path)
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Error downloading media")
|
||||
return None
|
||||
|
||||
async def _upload_media_ws(
|
||||
self, client: Any, file_path: str,
|
||||
) -> "tuple[str, str] | tuple[None, None]":
|
||||
"""Upload a local file to WeCom via WebSocket 3-step protocol (base64).
|
||||
|
||||
Uses the WeCom WebSocket upload commands directly via
|
||||
``client._ws_manager.send_reply()``:
|
||||
|
||||
``aibot_upload_media_init`` → upload_id
|
||||
``aibot_upload_media_chunk`` × N (≤512 KB raw per chunk, base64)
|
||||
``aibot_upload_media_finish`` → media_id
|
||||
|
||||
Returns (media_id, media_type) on success, (None, None) on failure.
|
||||
"""
|
||||
from wecom_aibot_sdk.utils import generate_req_id as _gen_req_id
|
||||
|
||||
try:
|
||||
fname = os.path.basename(file_path)
|
||||
media_type = _guess_wecom_media_type(fname)
|
||||
|
||||
# Read file size and data in a thread to avoid blocking the event loop
|
||||
def _read_file():
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size > WECOM_UPLOAD_MAX_BYTES:
|
||||
raise ValueError(
|
||||
f"File too large: {file_size} bytes (max {WECOM_UPLOAD_MAX_BYTES})"
|
||||
)
|
||||
with open(file_path, "rb") as f:
|
||||
return file_size, f.read()
|
||||
|
||||
file_size, data = await asyncio.to_thread(_read_file)
|
||||
# MD5 is used for file integrity only, not cryptographic security
|
||||
md5_hash = hashlib.md5(data).hexdigest()
|
||||
|
||||
chunk_size = 512 * 1024 # 512 KB raw (before base64)
|
||||
mv = memoryview(data)
|
||||
chunk_list = [bytes(mv[i : i + chunk_size]) for i in range(0, file_size, chunk_size)]
|
||||
n_chunks = len(chunk_list)
|
||||
del mv, data
|
||||
|
||||
# Step 1: init
|
||||
req_id = _gen_req_id("upload_init")
|
||||
resp = await client._ws_manager.send_reply(req_id, {
|
||||
"type": media_type,
|
||||
"filename": fname,
|
||||
"total_size": file_size,
|
||||
"total_chunks": n_chunks,
|
||||
"md5": md5_hash,
|
||||
}, "aibot_upload_media_init")
|
||||
if resp.errcode != 0:
|
||||
self.logger.warning("upload init failed ({}): {}", resp.errcode, resp.errmsg)
|
||||
return None, None
|
||||
upload_id = resp.body.get("upload_id") if resp.body else None
|
||||
if not upload_id:
|
||||
self.logger.warning("upload init: no upload_id in response")
|
||||
return None, None
|
||||
|
||||
# Step 2: send chunks
|
||||
for i, chunk in enumerate(chunk_list):
|
||||
req_id = _gen_req_id("upload_chunk")
|
||||
resp = await client._ws_manager.send_reply(req_id, {
|
||||
"upload_id": upload_id,
|
||||
"chunk_index": i,
|
||||
"base64_data": base64.b64encode(chunk).decode(),
|
||||
}, "aibot_upload_media_chunk")
|
||||
if resp.errcode != 0:
|
||||
self.logger.warning("upload chunk {} failed ({}): {}", i, resp.errcode, resp.errmsg)
|
||||
return None, None
|
||||
|
||||
# Step 3: finish
|
||||
req_id = _gen_req_id("upload_finish")
|
||||
resp = await client._ws_manager.send_reply(req_id, {
|
||||
"upload_id": upload_id,
|
||||
}, "aibot_upload_media_finish")
|
||||
if resp.errcode != 0:
|
||||
self.logger.warning("upload finish failed ({}): {}", resp.errcode, resp.errmsg)
|
||||
return None, None
|
||||
|
||||
media_id = resp.body.get("media_id") if resp.body else None
|
||||
if not media_id:
|
||||
self.logger.warning("upload finish: no media_id in response body={}", resp.body)
|
||||
return None, None
|
||||
|
||||
suffix = "..." if len(media_id) > 16 else ""
|
||||
self.logger.debug("uploaded {} ({}) → media_id={}", fname, media_type, media_id[:16] + suffix)
|
||||
return media_id, media_type
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.warning("upload skipped for {}: {}", file_path, e)
|
||||
return None, None
|
||||
except Exception:
|
||||
self.logger.exception("_upload_media_ws error for {}", file_path)
|
||||
return None, None
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
"""Send a message through WeCom."""
|
||||
if not self._client:
|
||||
self.logger.warning("client not initialized")
|
||||
return
|
||||
|
||||
try:
|
||||
content = (msg.content or "").strip()
|
||||
is_progress = bool(msg.metadata.get("_progress"))
|
||||
|
||||
# Get the stored frame for this chat
|
||||
frame = self._chat_frames.get(msg.chat_id)
|
||||
|
||||
# Send media files via WebSocket upload
|
||||
for file_path in msg.media or []:
|
||||
if not os.path.isfile(file_path):
|
||||
self.logger.warning("media file not found: {}", file_path)
|
||||
continue
|
||||
media_id, media_type = await self._upload_media_ws(self._client, file_path)
|
||||
if media_id:
|
||||
if frame:
|
||||
await self._client.reply(frame, {
|
||||
"msgtype": media_type,
|
||||
media_type: {"media_id": media_id},
|
||||
})
|
||||
else:
|
||||
await self._client.send_message(msg.chat_id, {
|
||||
"msgtype": media_type,
|
||||
media_type: {"media_id": media_id},
|
||||
})
|
||||
self.logger.debug("sent {} → {}", media_type, msg.chat_id)
|
||||
else:
|
||||
content += f"\n[file upload failed: {os.path.basename(file_path)}]"
|
||||
|
||||
if not content:
|
||||
return
|
||||
|
||||
if frame:
|
||||
# Both progress and final messages must use reply_stream (cmd="aibot_respond_msg").
|
||||
# The plain reply() uses cmd="reply" which does not support "text" msgtype
|
||||
# and causes errcode=40008 from WeCom API.
|
||||
stream_id = self._generate_req_id("stream")
|
||||
await self._client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
content,
|
||||
finish=not is_progress,
|
||||
)
|
||||
self.logger.debug(
|
||||
"{} sent to {}",
|
||||
"progress" if is_progress else "message",
|
||||
msg.chat_id,
|
||||
)
|
||||
else:
|
||||
# No frame (e.g. cron push): proactive send only supports markdown
|
||||
await self._client.send_message(msg.chat_id, {
|
||||
"msgtype": "markdown",
|
||||
"markdown": {"content": content},
|
||||
})
|
||||
self.logger.info("proactive send to {}", msg.chat_id)
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Error sending message to chat_id={}", msg.chat_id)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,682 @@
|
||||
"""WhatsApp channel implementation using neonize."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import mimetypes
|
||||
import re
|
||||
import secrets
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, NamedTuple
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from src.channels.bus.events import OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.base import BaseChannel
|
||||
from src.channels.utils import get_media_dir, get_runtime_subdir
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class WhatsAppConfig(BaseModel):
|
||||
"""WhatsApp channel configuration."""
|
||||
|
||||
enabled: bool = False
|
||||
allow_from: list[str] = Field(default_factory=list)
|
||||
group_policy: Literal["open", "mention"] = "open"
|
||||
database_path: str = ""
|
||||
lid_mappings: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class _NeonizeAPI(NamedTuple):
|
||||
NewAClient: Any
|
||||
ConnectedEv: Any
|
||||
DisconnectedEv: Any
|
||||
MessageEv: Any
|
||||
PairStatusEv: Any
|
||||
build_jid: Any
|
||||
|
||||
|
||||
class _MediaInfo(NamedTuple):
|
||||
kind: str
|
||||
message: Any
|
||||
mimetype: str
|
||||
filename: str
|
||||
is_voice: bool = False
|
||||
|
||||
|
||||
_NEONIZE_API: _NeonizeAPI | None = None
|
||||
_JID_RE = re.compile(r"^(?P<user>[^@]+)@(?P<server>[^@]+)$")
|
||||
_LEGACY_BRIDGE_CONFIG_FIELDS = ("bridgeUrl", "bridgeToken", "bridge_url", "bridge_token")
|
||||
|
||||
|
||||
def _default_database_path() -> Path:
|
||||
return get_runtime_subdir("whatsapp-auth") / "neonize.db"
|
||||
|
||||
|
||||
def _legacy_bridge_config_fields(config: dict[str, Any]) -> list[str]:
|
||||
return [field for field in _LEGACY_BRIDGE_CONFIG_FIELDS if field in config]
|
||||
|
||||
|
||||
def _load_neonize() -> _NeonizeAPI:
|
||||
global _NEONIZE_API
|
||||
if _NEONIZE_API is not None:
|
||||
return _NEONIZE_API
|
||||
|
||||
try:
|
||||
from neonize.aioze.client import NewAClient
|
||||
from neonize.aioze.events import ConnectedEv, DisconnectedEv, MessageEv, PairStatusEv
|
||||
from neonize.utils.jid import build_jid
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
'WhatsApp dependencies not installed. Run: pip install "vibe-trading-ai[whatsapp]"'
|
||||
) from exc
|
||||
|
||||
_NEONIZE_API = _NeonizeAPI(
|
||||
NewAClient=NewAClient,
|
||||
ConnectedEv=ConnectedEv,
|
||||
DisconnectedEv=DisconnectedEv,
|
||||
MessageEv=MessageEv,
|
||||
PairStatusEv=PairStatusEv,
|
||||
build_jid=build_jid,
|
||||
)
|
||||
return _NEONIZE_API
|
||||
|
||||
|
||||
def _has_field(message: Any, name: str) -> bool:
|
||||
if message is None:
|
||||
return False
|
||||
|
||||
has_field = getattr(message, "HasField", None)
|
||||
if callable(has_field):
|
||||
try:
|
||||
return bool(has_field(name))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
list_fields = getattr(message, "ListFields", None)
|
||||
if callable(list_fields):
|
||||
try:
|
||||
return any(getattr(field, "name", "") == name for field, _ in list_fields())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
value = getattr(message, name, None)
|
||||
return value is not None and value != "" and value != b""
|
||||
|
||||
|
||||
def _message_field(message: Any, *names: str) -> Any:
|
||||
for name in names:
|
||||
if _has_field(message, name):
|
||||
return getattr(message, name)
|
||||
return None
|
||||
|
||||
|
||||
def _safe_attr(obj: Any, name: str, default: Any = None) -> Any:
|
||||
if obj is None:
|
||||
return default
|
||||
return getattr(obj, name, default)
|
||||
|
||||
|
||||
def _jid_to_string(jid: Any) -> str:
|
||||
if jid is None:
|
||||
return ""
|
||||
if isinstance(jid, str):
|
||||
return jid.strip()
|
||||
if bool(_safe_attr(jid, "IsEmpty", False)):
|
||||
return ""
|
||||
|
||||
user = str(_safe_attr(jid, "User", "") or "").strip()
|
||||
server = str(_safe_attr(jid, "Server", "") or "").strip()
|
||||
if user and server:
|
||||
return f"{user}@{server}"
|
||||
return server or user
|
||||
|
||||
|
||||
def _normalize_jid(raw: Any) -> str:
|
||||
jid = _jid_to_string(raw).strip()
|
||||
if not jid:
|
||||
return ""
|
||||
if jid.endswith("@lid.whatsapp.net"):
|
||||
return jid[: -len(".whatsapp.net")]
|
||||
return jid
|
||||
|
||||
|
||||
def _bare_jid(raw: Any) -> str:
|
||||
jid = _normalize_jid(raw)
|
||||
if "@" not in jid:
|
||||
return jid
|
||||
return jid.split("@", 1)[0].split(":", 1)[0]
|
||||
|
||||
|
||||
def _classify_sender_ids(jids: list[Any]) -> tuple[str, str]:
|
||||
phone_id = ""
|
||||
lid_id = ""
|
||||
|
||||
for raw in jids:
|
||||
jid = _normalize_jid(raw)
|
||||
if not jid:
|
||||
continue
|
||||
match = _JID_RE.match(jid)
|
||||
if match:
|
||||
user = match.group("user").split(":", 1)[0]
|
||||
server = match.group("server")
|
||||
if server in {"s.whatsapp.net", "c.us"}:
|
||||
phone_id = phone_id or user
|
||||
elif server in {"lid", "lid.whatsapp.net"}:
|
||||
lid_id = lid_id or user
|
||||
continue
|
||||
|
||||
if not phone_id:
|
||||
phone_id = jid
|
||||
|
||||
return phone_id, lid_id
|
||||
|
||||
|
||||
def _context_infos(message: Any) -> list[Any]:
|
||||
infos: list[Any] = []
|
||||
for container in (
|
||||
message,
|
||||
_message_field(message, "extendedTextMessage"),
|
||||
_message_field(message, "imageMessage"),
|
||||
_message_field(message, "videoMessage"),
|
||||
_message_field(message, "audioMessage"),
|
||||
_message_field(message, "documentMessage"),
|
||||
_message_field(message, "stickerMessage"),
|
||||
):
|
||||
context = _message_field(container, "contextInfo")
|
||||
if context is not None:
|
||||
infos.append(context)
|
||||
return infos
|
||||
|
||||
|
||||
def _message_text(message: Any) -> str:
|
||||
conversation = str(_safe_attr(message, "conversation", "") or "").strip()
|
||||
if conversation:
|
||||
return conversation
|
||||
|
||||
extended = _message_field(message, "extendedTextMessage")
|
||||
text = str(_safe_attr(extended, "text", "") or "").strip()
|
||||
if text:
|
||||
return text
|
||||
|
||||
for field_name in ("imageMessage", "videoMessage", "documentMessage", "stickerMessage"):
|
||||
media_message = _message_field(message, field_name)
|
||||
caption = str(_safe_attr(media_message, "caption", "") or "").strip()
|
||||
if caption:
|
||||
return caption
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def _media_message(message: Any) -> _MediaInfo | None:
|
||||
image = _message_field(message, "imageMessage")
|
||||
if image is not None:
|
||||
return _MediaInfo(
|
||||
kind="image",
|
||||
message=image,
|
||||
mimetype=str(_safe_attr(image, "mimetype", "") or "image/jpeg"),
|
||||
filename=str(_safe_attr(image, "fileName", "") or ""),
|
||||
)
|
||||
|
||||
video = _message_field(message, "videoMessage")
|
||||
if video is not None:
|
||||
return _MediaInfo(
|
||||
kind="video",
|
||||
message=video,
|
||||
mimetype=str(_safe_attr(video, "mimetype", "") or "video/mp4"),
|
||||
filename=str(_safe_attr(video, "fileName", "") or ""),
|
||||
)
|
||||
|
||||
audio = _message_field(message, "audioMessage")
|
||||
if audio is not None:
|
||||
return _MediaInfo(
|
||||
kind="audio",
|
||||
message=audio,
|
||||
mimetype=str(_safe_attr(audio, "mimetype", "") or "audio/ogg"),
|
||||
filename=str(_safe_attr(audio, "fileName", "") or ""),
|
||||
is_voice=bool(_safe_attr(audio, "PTT", False) or _safe_attr(audio, "ptt", False)),
|
||||
)
|
||||
|
||||
document = _message_field(message, "documentMessage")
|
||||
if document is not None:
|
||||
return _MediaInfo(
|
||||
kind="file",
|
||||
message=document,
|
||||
mimetype=str(_safe_attr(document, "mimetype", "") or "application/octet-stream"),
|
||||
filename=str(
|
||||
_safe_attr(document, "fileName", "")
|
||||
or _safe_attr(document, "title", "")
|
||||
or ""
|
||||
),
|
||||
)
|
||||
|
||||
sticker = _message_field(message, "stickerMessage")
|
||||
if sticker is not None:
|
||||
return _MediaInfo(
|
||||
kind="sticker",
|
||||
message=sticker,
|
||||
mimetype=str(_safe_attr(sticker, "mimetype", "") or "image/webp"),
|
||||
filename=str(_safe_attr(sticker, "fileName", "") or ""),
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class WhatsAppChannel(BaseChannel):
|
||||
"""WhatsApp channel using neonize's async WhatsApp client."""
|
||||
|
||||
name = "whatsapp"
|
||||
display_name = "WhatsApp"
|
||||
|
||||
@classmethod
|
||||
def default_config(cls) -> dict[str, Any]:
|
||||
return WhatsAppConfig().model_dump(by_alias=True)
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
legacy_bridge_fields = _legacy_bridge_config_fields(config) if isinstance(config, dict) else []
|
||||
if isinstance(config, dict):
|
||||
config = WhatsAppConfig.model_validate(config)
|
||||
super().__init__(config, bus)
|
||||
if legacy_bridge_fields:
|
||||
self.logger.warning(
|
||||
"Ignoring deprecated WhatsApp bridge config fields: {}. "
|
||||
"Run 'vibe-trading channels login whatsapp' to create a neonize session.",
|
||||
", ".join(legacy_bridge_fields),
|
||||
)
|
||||
self._client: Any | None = None
|
||||
self._connected = False
|
||||
self._processed_message_ids: OrderedDict[str, None] = OrderedDict()
|
||||
self._lid_to_phone = self._load_lid_mappings()
|
||||
self._self_jids: set[str] = set()
|
||||
self._started_at = 0.0
|
||||
|
||||
def _database_path(self) -> Path:
|
||||
configured = self.config.database_path.strip()
|
||||
return Path(configured).expanduser() if configured else _default_database_path()
|
||||
|
||||
def _load_lid_mappings(self) -> dict[str, str]:
|
||||
mapping: dict[str, str] = {}
|
||||
for lid, phone in self.config.lid_mappings.items():
|
||||
phone_text = str(phone).strip()
|
||||
if phone_text:
|
||||
mapping[str(lid).strip()] = phone_text
|
||||
return mapping
|
||||
|
||||
def _new_client(self) -> Any:
|
||||
api = _load_neonize()
|
||||
db_path = self._database_path()
|
||||
db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return api.NewAClient(str(db_path))
|
||||
|
||||
async def login(self, force: bool = False) -> bool:
|
||||
db_path = self._database_path()
|
||||
if force:
|
||||
self._reset_database(db_path)
|
||||
|
||||
client = self._new_client()
|
||||
login_result = asyncio.get_running_loop().create_future()
|
||||
self._register_handlers(client, login_result=login_result, handle_messages=False)
|
||||
|
||||
try:
|
||||
self.logger.info("Starting WhatsApp login with neonize...")
|
||||
connect_task = await client.connect()
|
||||
self._fail_login_on_connect_task_done(connect_task, login_result)
|
||||
await login_result
|
||||
self.logger.info("WhatsApp login complete")
|
||||
return True
|
||||
except Exception as exc:
|
||||
self.logger.error("WhatsApp login failed: {}", exc)
|
||||
return False
|
||||
finally:
|
||||
with suppress(Exception):
|
||||
await client.stop()
|
||||
|
||||
async def start(self) -> None:
|
||||
self._running = True
|
||||
self._started_at = time.time()
|
||||
client = self._new_client()
|
||||
self._client = client
|
||||
self._register_handlers(client, handle_messages=True)
|
||||
|
||||
try:
|
||||
self.logger.info("Connecting WhatsApp channel with neonize...")
|
||||
await client.connect()
|
||||
await client.idle()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
finally:
|
||||
self._running = False
|
||||
self._connected = False
|
||||
if self._client is client:
|
||||
self._client = None
|
||||
with suppress(Exception):
|
||||
await client.stop()
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._running = False
|
||||
self._connected = False
|
||||
client = self._client
|
||||
self._client = None
|
||||
if client is not None:
|
||||
await client.stop()
|
||||
|
||||
@staticmethod
|
||||
def _fail_login_on_connect_task_done(
|
||||
connect_task: asyncio.Task[Any] | None,
|
||||
login_result: asyncio.Future[None],
|
||||
) -> None:
|
||||
if connect_task is None:
|
||||
return
|
||||
|
||||
def _on_done(task: asyncio.Task[Any]) -> None:
|
||||
try:
|
||||
exc = task.exception()
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
if login_result.done():
|
||||
return
|
||||
if exc is not None:
|
||||
login_result.set_exception(exc)
|
||||
else:
|
||||
login_result.set_exception(
|
||||
RuntimeError("WhatsApp connection ended before login completed")
|
||||
)
|
||||
|
||||
connect_task.add_done_callback(_on_done)
|
||||
|
||||
async def send(self, msg: OutboundMessage) -> None:
|
||||
client = self._client
|
||||
if client is None or not self._connected:
|
||||
raise RuntimeError("WhatsApp channel is not connected")
|
||||
|
||||
to = self._build_jid(msg.chat_id)
|
||||
if msg.content:
|
||||
await client.send_message(to, msg.content)
|
||||
|
||||
for media_path in msg.media or []:
|
||||
await self._send_media(client, to, media_path)
|
||||
|
||||
def _build_jid(self, raw: str) -> Any:
|
||||
api = _load_neonize()
|
||||
target = raw.strip()
|
||||
match = _JID_RE.match(_normalize_jid(target))
|
||||
if not match:
|
||||
return api.build_jid(target)
|
||||
|
||||
user = match.group("user").split(":", 1)[0]
|
||||
server = match.group("server")
|
||||
return api.build_jid(user, server)
|
||||
|
||||
async def _send_media(self, client: Any, to: Any, media_path: str) -> None:
|
||||
path = str(Path(media_path).expanduser())
|
||||
mime, _ = mimetypes.guess_type(path)
|
||||
mimetype = mime or "application/octet-stream"
|
||||
if mimetype.startswith("image/"):
|
||||
await client.send_image(to, path)
|
||||
elif mimetype.startswith("video/"):
|
||||
await client.send_video(to, path)
|
||||
elif mimetype.startswith("audio/"):
|
||||
await client.send_audio(to, path)
|
||||
else:
|
||||
await client.send_document(
|
||||
to,
|
||||
path,
|
||||
filename=Path(path).name,
|
||||
mimetype=mimetype,
|
||||
)
|
||||
|
||||
def _register_handlers(
|
||||
self,
|
||||
client: Any,
|
||||
*,
|
||||
login_result: asyncio.Future[None] | None = None,
|
||||
handle_messages: bool,
|
||||
) -> None:
|
||||
api = _load_neonize()
|
||||
|
||||
@client.qr
|
||||
async def _on_qr(_: Any, qr_data: bytes) -> None:
|
||||
import segno
|
||||
|
||||
self.logger.info("Scan the WhatsApp QR code with Linked Devices")
|
||||
segno.make_qr(qr_data).terminal(compact=True)
|
||||
|
||||
@client.event(api.ConnectedEv)
|
||||
async def _on_connected(current_client: Any, _: Any) -> None:
|
||||
self._connected = True
|
||||
try:
|
||||
await self._remember_self_jids(current_client)
|
||||
except Exception as exc:
|
||||
if login_result is not None and not login_result.done():
|
||||
login_result.set_exception(exc)
|
||||
raise
|
||||
if login_result is not None and not login_result.done():
|
||||
login_result.set_result(None)
|
||||
self.logger.info("WhatsApp connected")
|
||||
|
||||
@client.event(api.DisconnectedEv)
|
||||
async def _on_disconnected(_: Any, event: Any) -> None:
|
||||
self._connected = False
|
||||
if login_result is not None and not login_result.done():
|
||||
login_result.set_exception(
|
||||
RuntimeError(f"WhatsApp disconnected before login completed: {event}")
|
||||
)
|
||||
self.logger.warning("WhatsApp disconnected: {}", event)
|
||||
|
||||
@client.event(api.PairStatusEv)
|
||||
async def _on_pair_status(_: Any, event: Any) -> None:
|
||||
error = str(_safe_attr(event, "Error", "") or "")
|
||||
if error:
|
||||
exc = RuntimeError(f"WhatsApp pair status error: {error}")
|
||||
if login_result is not None and not login_result.done():
|
||||
login_result.set_exception(exc)
|
||||
raise exc
|
||||
self.logger.info("WhatsApp pair status: {}", event)
|
||||
|
||||
if not handle_messages:
|
||||
return
|
||||
|
||||
@client.event(api.MessageEv)
|
||||
async def _on_message(current_client: Any, event: Any) -> None:
|
||||
try:
|
||||
await self._handle_neonize_message(current_client, event)
|
||||
except Exception:
|
||||
self.logger.exception("Error handling WhatsApp message")
|
||||
raise
|
||||
|
||||
async def _remember_self_jids(self, client: Any) -> None:
|
||||
device = _safe_attr(client, "me")
|
||||
if device is None:
|
||||
device = await client.get_me()
|
||||
|
||||
for attr in ("JID", "LID"):
|
||||
jid = _normalize_jid(_safe_attr(device, attr))
|
||||
if jid:
|
||||
self._self_jids.add(jid)
|
||||
self._self_jids.add(_bare_jid(jid))
|
||||
|
||||
async def _handle_neonize_message(self, client: Any, event: Any) -> None:
|
||||
info = _safe_attr(event, "Info")
|
||||
message = _safe_attr(event, "Message")
|
||||
source = _safe_attr(info, "MessageSource")
|
||||
if info is None or message is None or source is None:
|
||||
raise ValueError("WhatsApp MessageEv is missing Info, Message, or MessageSource")
|
||||
|
||||
if bool(_safe_attr(source, "IsFromMe", False)):
|
||||
return
|
||||
|
||||
chat_jid = _normalize_jid(_safe_attr(source, "Chat"))
|
||||
if not chat_jid:
|
||||
raise ValueError("WhatsApp message has no chat JID")
|
||||
if chat_jid == "status@broadcast":
|
||||
return
|
||||
|
||||
timestamp = float(_safe_attr(info, "Timestamp", 0) or 0)
|
||||
if self._started_at and timestamp and timestamp < self._started_at:
|
||||
return
|
||||
|
||||
is_group = bool(_safe_attr(source, "IsGroup", False))
|
||||
if is_group and self.config.group_policy == "mention":
|
||||
if not self._is_addressed_to_bot(message):
|
||||
return
|
||||
|
||||
message_id = str(_safe_attr(info, "ID", "") or "")
|
||||
if message_id:
|
||||
if message_id in self._processed_message_ids:
|
||||
return
|
||||
self._processed_message_ids[message_id] = None
|
||||
while len(self._processed_message_ids) > 1000:
|
||||
self._processed_message_ids.popitem(last=False)
|
||||
|
||||
participant_jid = _normalize_jid(_safe_attr(source, "Sender"))
|
||||
sender_alt_jid = _normalize_jid(_safe_attr(source, "SenderAlt"))
|
||||
sender_candidates = [sender_alt_jid, participant_jid]
|
||||
if not is_group:
|
||||
sender_candidates.append(chat_jid)
|
||||
|
||||
phone_id, lid_id = _classify_sender_ids(sender_candidates)
|
||||
if phone_id and lid_id:
|
||||
self._lid_to_phone[lid_id] = phone_id
|
||||
|
||||
sender_id = phone_id or self._lid_to_phone.get(lid_id, "") or lid_id
|
||||
if not sender_id:
|
||||
raise ValueError("WhatsApp message has no resolvable sender ID")
|
||||
metadata = {
|
||||
"message_id": message_id or None,
|
||||
"timestamp": int(timestamp) if timestamp else None,
|
||||
"is_group": is_group,
|
||||
"is_forwarded": self._is_forwarded(message),
|
||||
"participant": participant_jid or None,
|
||||
"sender_alt": sender_alt_jid or None,
|
||||
"lid": lid_id or None,
|
||||
"phone": phone_id or None,
|
||||
"is_reply_to_bot": self._is_reply_to_bot(message),
|
||||
}
|
||||
if not self.is_allowed(sender_id):
|
||||
self.logger.info(
|
||||
"Passing unauthorized WhatsApp sender {} to pairing flow "
|
||||
"(phone={}, lid={}, chat={})",
|
||||
sender_id,
|
||||
phone_id or "",
|
||||
lid_id or "",
|
||||
chat_jid,
|
||||
)
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_jid,
|
||||
content=_message_text(message),
|
||||
media=[],
|
||||
metadata=metadata,
|
||||
is_dm=not is_group,
|
||||
)
|
||||
return
|
||||
|
||||
text = _message_text(message)
|
||||
media_paths: list[str] = []
|
||||
media = _media_message(message)
|
||||
if media is not None:
|
||||
path = await self._download_media(client, event, media)
|
||||
if media.kind == "audio" and media.is_voice:
|
||||
transcription = await self.transcribe_audio(path)
|
||||
if transcription:
|
||||
text = transcription
|
||||
else:
|
||||
media_paths.append(path)
|
||||
text = self._append_media_tag(text, "audio", path)
|
||||
else:
|
||||
media_paths.append(path)
|
||||
text = self._append_media_tag(text, media.kind, path)
|
||||
|
||||
if not text and not media_paths:
|
||||
return
|
||||
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_jid,
|
||||
content=text,
|
||||
media=media_paths,
|
||||
metadata=metadata,
|
||||
is_dm=not is_group,
|
||||
)
|
||||
|
||||
def _is_addressed_to_bot(self, message: Any) -> bool:
|
||||
return self._was_mentioned(message) or self._is_reply_to_bot(message)
|
||||
|
||||
def _was_mentioned(self, message: Any) -> bool:
|
||||
if not self._self_jids:
|
||||
return False
|
||||
for context in _context_infos(message):
|
||||
mentioned = (
|
||||
_safe_attr(context, "mentionedJID")
|
||||
or _safe_attr(context, "mentionedJid")
|
||||
or _safe_attr(context, "mentioned_jid")
|
||||
or []
|
||||
)
|
||||
for jid in mentioned:
|
||||
normalized = _normalize_jid(jid)
|
||||
if normalized in self._self_jids or _bare_jid(normalized) in self._self_jids:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _is_reply_to_bot(self, message: Any) -> bool:
|
||||
if not self._self_jids:
|
||||
return False
|
||||
for context in _context_infos(message):
|
||||
participant = _normalize_jid(
|
||||
_safe_attr(context, "participant")
|
||||
or _safe_attr(context, "Participant")
|
||||
or ""
|
||||
)
|
||||
if participant in self._self_jids or _bare_jid(participant) in self._self_jids:
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_forwarded(message: Any) -> bool:
|
||||
for context in _context_infos(message):
|
||||
if bool(_safe_attr(context, "isForwarded", False)):
|
||||
return True
|
||||
if int(_safe_attr(context, "forwardingScore", 0) or 0) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _download_media(self, client: Any, event: Any, media: _MediaInfo) -> str:
|
||||
info = _safe_attr(event, "Info")
|
||||
message_id = str(_safe_attr(info, "ID", "") or "")
|
||||
path = self._media_path(message_id, media)
|
||||
await client.download_any(_safe_attr(event, "Message"), str(path))
|
||||
return str(path)
|
||||
|
||||
def _media_path(self, message_id: str, media: _MediaInfo) -> Path:
|
||||
media_dir = get_media_dir("whatsapp")
|
||||
safe_id = re.sub(r"[^A-Za-z0-9_.-]+", "_", message_id or str(int(time.time())))
|
||||
filename = Path(media.filename).name if media.filename else ""
|
||||
suffix = Path(filename).suffix if filename else ""
|
||||
if not suffix:
|
||||
suffix = mimetypes.guess_extension(media.mimetype) or {
|
||||
"image": ".jpg",
|
||||
"video": ".mp4",
|
||||
"audio": ".ogg",
|
||||
"sticker": ".webp",
|
||||
}.get(media.kind, ".bin")
|
||||
return media_dir / f"wa_{safe_id}_{secrets.token_hex(4)}{suffix}"
|
||||
|
||||
@staticmethod
|
||||
def _append_media_tag(text: str, kind: str, path: str) -> str:
|
||||
label = kind if kind in {"image", "video", "audio", "sticker"} else "file"
|
||||
tag = f"[{label}: {path}]"
|
||||
return f"{text}\n{tag}" if text else tag
|
||||
|
||||
@staticmethod
|
||||
def _reset_database(path: Path) -> None:
|
||||
for candidate in (
|
||||
path,
|
||||
path.with_suffix(path.suffix + "-shm"),
|
||||
path.with_suffix(path.suffix + "-wal"),
|
||||
):
|
||||
if candidate.exists():
|
||||
candidate.unlink()
|
||||
@@ -0,0 +1 @@
|
||||
"""WebUI compatibility helpers for channel adapters."""
|
||||
@@ -0,0 +1,19 @@
|
||||
"""CLI app mention normalization for the WebSocket channel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def normalize_cli_app_mentions(content: Any, metadata: dict[str, Any] | None = None) -> list[str]:
|
||||
"""Normalize CLI app mentions into a stable list of names."""
|
||||
del metadata
|
||||
if content is None:
|
||||
return []
|
||||
if isinstance(content, str):
|
||||
raw_items = content.replace(",", " ").split()
|
||||
elif isinstance(content, (list, tuple, set)):
|
||||
raw_items = [str(item) for item in content]
|
||||
else:
|
||||
raw_items = [str(content)]
|
||||
return [item.lstrip("@").strip() for item in raw_items if item and item.strip()]
|
||||
@@ -0,0 +1,23 @@
|
||||
"""WebUI fork-chat compatibility for the WebSocket channel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
_HANDLE_WEBUI_FORK_CHAT = None
|
||||
|
||||
|
||||
async def handle_webui_fork_chat(*args: Any, **kwargs: Any) -> None:
|
||||
"""Return a structured unsupported response for fork-chat requests."""
|
||||
del kwargs
|
||||
if len(args) >= 2:
|
||||
channel, connection = args[0], args[1]
|
||||
send_event = getattr(channel, "_send_event", None)
|
||||
if callable(send_event):
|
||||
await send_event(
|
||||
connection,
|
||||
"error",
|
||||
detail="fork_chat is not available in this runtime",
|
||||
)
|
||||
return
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Gateway services used by the WebSocket channel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import secrets
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from src.config.paths import get_workspace_path
|
||||
from src.security.workspace_access import WorkspaceScopeError
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebSocketTokenIssuer:
|
||||
"""Short-lived one-time token issuer for WebSocket clients."""
|
||||
|
||||
_issued: dict[str, float] = field(default_factory=dict)
|
||||
|
||||
def issue_token(self, *, ttl_s: int = 300) -> tuple[str, int]:
|
||||
"""Issue a one-time token.
|
||||
|
||||
Args:
|
||||
ttl_s: Token lifetime in seconds.
|
||||
|
||||
Returns:
|
||||
``(token, expires_in_seconds)``.
|
||||
"""
|
||||
self._gc()
|
||||
token = secrets.token_urlsafe(32)
|
||||
self._issued[token] = time.time() + ttl_s
|
||||
return token, ttl_s
|
||||
|
||||
def take_issued_token_if_valid(self, token: str | None) -> bool:
|
||||
"""Consume a valid issued token exactly once."""
|
||||
if not token:
|
||||
return False
|
||||
self._gc()
|
||||
expires_at = self._issued.pop(token, None)
|
||||
return bool(expires_at and expires_at >= time.time())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all issued tokens."""
|
||||
self._issued.clear()
|
||||
|
||||
def _gc(self) -> None:
|
||||
now = time.time()
|
||||
for token, expires_at in list(self._issued.items()):
|
||||
if expires_at < now:
|
||||
self._issued.pop(token, None)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkspaceScope:
|
||||
"""Workspace scope snapshot carried in WebSocket session metadata."""
|
||||
|
||||
root: str
|
||||
restrict_to_workspace: bool = True
|
||||
|
||||
def payload(self) -> dict[str, Any]:
|
||||
"""Return the client-facing workspace scope."""
|
||||
return {
|
||||
"root": self.root,
|
||||
"restrict_to_workspace": self.restrict_to_workspace,
|
||||
}
|
||||
|
||||
def metadata(self) -> dict[str, Any]:
|
||||
"""Return the session metadata representation."""
|
||||
return self.payload()
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkspaceService:
|
||||
"""Resolve and persist per-chat workspace scope."""
|
||||
|
||||
workspace_path: Path = field(default_factory=get_workspace_path)
|
||||
default_restrict_to_workspace: bool = True
|
||||
_scopes: dict[str, WorkspaceScope] = field(default_factory=dict)
|
||||
|
||||
def _default_scope(self) -> WorkspaceScope:
|
||||
return WorkspaceScope(
|
||||
root=str(self.workspace_path.expanduser().resolve()),
|
||||
restrict_to_workspace=self.default_restrict_to_workspace,
|
||||
)
|
||||
|
||||
def scope_for_new_chat(self, envelope: dict[str, Any], *, controls_available: bool) -> WorkspaceScope:
|
||||
"""Build a scope for a new chat envelope."""
|
||||
del controls_available
|
||||
return self._scope_from_envelope(envelope) or self._default_scope()
|
||||
|
||||
def scope_for_set_request(
|
||||
self,
|
||||
envelope: dict[str, Any],
|
||||
*,
|
||||
chat_id: str,
|
||||
chat_running: bool,
|
||||
controls_available: bool,
|
||||
) -> WorkspaceScope:
|
||||
"""Build a scope update from a client request."""
|
||||
del chat_running, controls_available
|
||||
return self._scope_from_envelope(envelope) or self._scopes.get(chat_id) or self._default_scope()
|
||||
|
||||
def scope_for_message(
|
||||
self,
|
||||
envelope: dict[str, Any],
|
||||
*,
|
||||
chat_id: str,
|
||||
chat_running: bool,
|
||||
controls_available: bool,
|
||||
) -> WorkspaceScope:
|
||||
"""Return the active scope for a message."""
|
||||
del chat_running, controls_available
|
||||
return self._scope_from_envelope(envelope) or self._scopes.get(chat_id) or self._default_scope()
|
||||
|
||||
def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
|
||||
"""Persist an in-memory scope for the current gateway process."""
|
||||
self._scopes[chat_id] = scope
|
||||
|
||||
def _scope_from_envelope(self, envelope: dict[str, Any]) -> WorkspaceScope | None:
|
||||
raw = envelope.get("workspace_scope")
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
root = raw.get("root")
|
||||
if not isinstance(root, str) or not root.strip():
|
||||
return None
|
||||
resolved = Path(root).expanduser().resolve()
|
||||
default_root = self.workspace_path.expanduser().resolve()
|
||||
if self.default_restrict_to_workspace and not _is_relative_to(resolved, default_root):
|
||||
raise WorkspaceScopeError(f"workspace root must stay under {default_root}")
|
||||
restrict = raw.get("restrict_to_workspace", self.default_restrict_to_workspace)
|
||||
return WorkspaceScope(root=str(resolved), restrict_to_workspace=bool(restrict))
|
||||
|
||||
|
||||
class SimpleHttpRouter:
|
||||
"""Small HTTP fallback router for WebSocket server requests."""
|
||||
|
||||
def workspace_controls_available(self, connection: Any) -> bool:
|
||||
"""Return whether workspace controls may be shown to this connection."""
|
||||
del connection
|
||||
return True
|
||||
|
||||
async def dispatch(self, connection: Any, request: Any) -> Any:
|
||||
"""Return a compact JSON 404 for non-WebSocket HTTP requests."""
|
||||
del request
|
||||
return connection.respond(
|
||||
404,
|
||||
json.dumps({"detail": "not found"}, ensure_ascii=False),
|
||||
)
|
||||
|
||||
|
||||
class MediaService:
|
||||
"""Media staging helpers for WebSocket payloads."""
|
||||
|
||||
def rewrite_local_markdown_images(self, text: str) -> str:
|
||||
"""Return text unchanged until an HTTP media server is configured."""
|
||||
return text
|
||||
|
||||
def sign_or_stage_media_path(self, path: Path) -> dict[str, str] | None:
|
||||
"""Return metadata for a local media path when it exists."""
|
||||
try:
|
||||
resolved = path.expanduser().resolve()
|
||||
except OSError:
|
||||
return None
|
||||
if not resolved.exists() or not resolved.is_file():
|
||||
return None
|
||||
return {"name": resolved.name, "path": str(resolved)}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscriptService:
|
||||
"""In-process transcript event bridge for the WebSocket channel."""
|
||||
|
||||
events: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
def client_turn_metadata(self, turn_id: Any) -> dict[str, Any]:
|
||||
"""Return metadata derived from a client turn id."""
|
||||
return {"turn_id": turn_id} if isinstance(turn_id, str) and turn_id else {}
|
||||
|
||||
def append_user_message(
|
||||
self,
|
||||
chat_id: str,
|
||||
content: str,
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
media_paths: list[str] | None = None,
|
||||
cli_apps: list[str] | None = None,
|
||||
mcp_presets: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Record a user-message transcript event for the current process."""
|
||||
self.events.append(
|
||||
{
|
||||
"phase": "user",
|
||||
"chat_id": chat_id,
|
||||
"content": content,
|
||||
"metadata": metadata or {},
|
||||
"media_paths": media_paths or [],
|
||||
"cli_apps": cli_apps or [],
|
||||
"mcp_presets": mcp_presets or [],
|
||||
}
|
||||
)
|
||||
|
||||
def prepare_and_append(
|
||||
self,
|
||||
chat_id: str,
|
||||
payload: dict[str, Any],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
phase: str = "",
|
||||
include_source: bool = False,
|
||||
transcript_overrides: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Record a generated transcript event for the current process."""
|
||||
self.events.append(
|
||||
{
|
||||
"phase": phase,
|
||||
"chat_id": chat_id,
|
||||
"payload": payload,
|
||||
"metadata": metadata or {},
|
||||
"include_source": include_source,
|
||||
"transcript_overrides": transcript_overrides or {},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class GatewaySessionManagerAdapter:
|
||||
"""Adapt Vibe-Trading SessionService to the WebSocket gateway contract."""
|
||||
|
||||
def __init__(self, service: Any) -> None:
|
||||
self._service = service
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._service, name)
|
||||
|
||||
def read_session_file(self, session_id: str) -> dict[str, Any]:
|
||||
"""Return a JSON-like session row for WebSocket hydration."""
|
||||
getter = getattr(self._service, "get_session", None)
|
||||
if not callable(getter):
|
||||
return {"metadata": {}}
|
||||
session = getter(session_id)
|
||||
if session is None and ":" in session_id:
|
||||
session = getter(session_id.split(":", 1)[1])
|
||||
if session is None:
|
||||
return {"metadata": {}}
|
||||
data = session.to_dict() if hasattr(session, "to_dict") else dict(getattr(session, "__dict__", {}))
|
||||
config = data.get("config") if isinstance(data.get("config"), dict) else {}
|
||||
metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else config.get("metadata", {})
|
||||
data["metadata"] = metadata if isinstance(metadata, dict) else {}
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
class GatewayServices:
|
||||
"""Service bundle required by :class:`src.channels.websocket.WebSocketChannel`."""
|
||||
|
||||
session_manager: Any = None
|
||||
cron_service: Any = None
|
||||
runtime_model_name: Any = None
|
||||
cron_pending_job_ids: Any = None
|
||||
static_dist_path: Any = None
|
||||
workspace_path: Path | None = None
|
||||
default_restrict_to_workspace: bool = True
|
||||
disabled_skills: set[str] = field(default_factory=set)
|
||||
logger: Any = None
|
||||
runtime_surface: str = "browser"
|
||||
runtime_capabilities: dict[str, Any] = field(default_factory=dict)
|
||||
http: SimpleHttpRouter = field(default_factory=SimpleHttpRouter)
|
||||
tokens: WebSocketTokenIssuer = field(default_factory=WebSocketTokenIssuer)
|
||||
media: MediaService = field(default_factory=MediaService)
|
||||
transcripts: TranscriptService = field(default_factory=TranscriptService)
|
||||
workspaces: WorkspaceService = field(default_factory=WorkspaceService)
|
||||
|
||||
|
||||
def build_gateway_services(**kwargs: Any) -> GatewayServices:
|
||||
"""Build a gateway service bundle from optional overrides."""
|
||||
allowed = {
|
||||
key: value
|
||||
for key, value in kwargs.items()
|
||||
if key in GatewayServices.__dataclass_fields__
|
||||
}
|
||||
session_manager = allowed.get("session_manager")
|
||||
if session_manager is not None and not hasattr(session_manager, "read_session_file"):
|
||||
allowed["session_manager"] = GatewaySessionManagerAdapter(session_manager)
|
||||
workspace_path = allowed.get("workspace_path")
|
||||
if workspace_path is not None and "workspaces" not in allowed:
|
||||
allowed["workspaces"] = WorkspaceService(
|
||||
workspace_path=Path(workspace_path),
|
||||
default_restrict_to_workspace=bool(
|
||||
allowed.get("default_restrict_to_workspace", True)
|
||||
),
|
||||
)
|
||||
return GatewayServices(**allowed)
|
||||
|
||||
|
||||
def _is_relative_to(path: Path, root: Path) -> bool:
|
||||
try:
|
||||
path.relative_to(root)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
@@ -0,0 +1,78 @@
|
||||
"""HTTP helpers for the WebSocket channel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
from src.channels.utils import validate_url_target
|
||||
|
||||
|
||||
_LONG_LIVED_TIMEOUT = 300 # seconds for long-lived HTTP connections
|
||||
|
||||
|
||||
def normalize_config_path(path: str | Path) -> str:
|
||||
"""Normalize an HTTP/WebSocket route path.
|
||||
|
||||
Args:
|
||||
path: Path-like route value.
|
||||
|
||||
Returns:
|
||||
A leading-slash route without a trailing slash, except for ``"/"``.
|
||||
"""
|
||||
value = str(path).strip() or "/"
|
||||
if not value.startswith("/"):
|
||||
value = f"/{value}"
|
||||
while "//" in value:
|
||||
value = value.replace("//", "/")
|
||||
if len(value) > 1:
|
||||
value = value.rstrip("/")
|
||||
return value or "/"
|
||||
|
||||
|
||||
def parse_request_path(path: str) -> tuple[str, dict[str, list[str]]]:
|
||||
"""Parse an HTTP request path into ``(route_path, query)``."""
|
||||
parts = urlsplit(path or "/")
|
||||
return normalize_config_path(parts.path or "/"), parse_qs(parts.query, keep_blank_values=True)
|
||||
|
||||
|
||||
def query_first(qs: dict[str, list[str]], key: str) -> str | None:
|
||||
"""Return the first value for *key* in a query string dict."""
|
||||
values = qs.get(key, [])
|
||||
return values[0] if values else None
|
||||
|
||||
|
||||
def parse_and_validate_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
"""Validate a URL target before a channel-side fetch."""
|
||||
return validate_url_target(url, allow_loopback=allow_loopback)
|
||||
|
||||
|
||||
async def read_uploaded_file(*args: Any, **kwargs: Any) -> bytes:
|
||||
"""Read bytes from a file-like upload object.
|
||||
|
||||
The helper accepts common async/sync file objects used by lightweight HTTP
|
||||
adapters. It is intentionally conservative and only returns raw bytes.
|
||||
"""
|
||||
del kwargs
|
||||
if not args:
|
||||
return b""
|
||||
obj = args[0]
|
||||
if isinstance(obj, bytes):
|
||||
return obj
|
||||
if isinstance(obj, bytearray):
|
||||
return bytes(obj)
|
||||
read = getattr(obj, "read", None)
|
||||
if read is None:
|
||||
raise TypeError("upload object does not provide read()")
|
||||
result = read()
|
||||
if asyncio.iscoroutine(result):
|
||||
result = await result
|
||||
if isinstance(result, str):
|
||||
return result.encode("utf-8")
|
||||
if isinstance(result, bytes):
|
||||
return result
|
||||
if isinstance(result, bytearray):
|
||||
return bytes(result)
|
||||
raise TypeError("upload read() did not return bytes")
|
||||
@@ -0,0 +1,19 @@
|
||||
"""MCP preset mention normalization for the WebSocket channel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def normalize_mcp_preset_mentions(content: Any, metadata: dict[str, Any] | None = None) -> list[str]:
|
||||
"""Normalize MCP preset mentions into a stable list of names."""
|
||||
del metadata
|
||||
if content is None:
|
||||
return []
|
||||
if isinstance(content, str):
|
||||
raw_items = content.replace(",", " ").split()
|
||||
elif isinstance(content, (list, tuple, set)):
|
||||
raw_items = [str(item) for item in content]
|
||||
else:
|
||||
raw_items = [str(content)]
|
||||
return [item.lstrip("@").strip() for item in raw_items if item and item.strip()]
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Transcription WebSocket compatibility helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def webui_transcription_event(*args: Any, **kwargs: Any) -> tuple[str, dict[str, Any]]:
|
||||
"""Return a structured response when transcription is not configured."""
|
||||
del args, kwargs
|
||||
return "error", {"detail": "audio transcription is not configured"}
|
||||
@@ -0,0 +1,7 @@
|
||||
"""WebSocket logging bridge for channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
websockets_server_logger = logging.getLogger("src.channelsui.websocket")
|
||||
@@ -68,4 +68,15 @@ def get_data_dir(config_path: Path | None = None) -> Path:
|
||||
"""
|
||||
data_dir = get_config_path(config_path).parent
|
||||
data_dir.mkdir(parents=True, exist_ok=True)
|
||||
return data_dir
|
||||
return data_dir
|
||||
|
||||
|
||||
def get_workspace_path() -> Path:
|
||||
"""Return the workspace path for channel state data.
|
||||
|
||||
For channel adapters that need to persist state (e.g. conversation
|
||||
references, auth tokens), this returns ``~/.vibe-trading/workspace``.
|
||||
"""
|
||||
p = get_runtime_root() / "workspace"
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
@@ -424,10 +424,26 @@ class MCPServerConfigOverride(ConfigBase):
|
||||
enabled_tools: list[str] | None = None
|
||||
|
||||
|
||||
class ChannelsConfig(ConfigBase):
|
||||
"""Top-level IM channel config.
|
||||
|
||||
Built-in adapters parse their own per-channel sections. This model keeps
|
||||
the operator file strict for global fields while allowing platform-specific
|
||||
channel keys such as ``telegram`` or ``feishu``.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(alias_generator=_to_camel, populate_by_name=True, extra="allow")
|
||||
|
||||
send_progress: bool = True
|
||||
send_tool_hints: bool = False
|
||||
send_max_retries: int = Field(default=2, ge=1, le=10)
|
||||
|
||||
|
||||
class AgentConfig(ConfigBase):
|
||||
"""Top-level structured agent config."""
|
||||
|
||||
mcp_servers: dict[str, MCPServerConfig] = Field(default_factory=dict)
|
||||
channels: ChannelsConfig = Field(default_factory=ChannelsConfig)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_live_broker_servers(self) -> "AgentConfig":
|
||||
@@ -483,3 +499,4 @@ class AgentConfigOverride(ConfigBase):
|
||||
)
|
||||
|
||||
mcp_servers: dict[str, MCPServerConfigOverride] = Field(default_factory=dict)
|
||||
channels: ChannelsConfig | None = None
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Crash recovery + position reconciliation (SPEC.md §7.5 component 5).
|
||||
|
||||
This is the hardest piece of the persistent runtime and has **no nanobot
|
||||
This is the hardest piece of the persistent runtime and has **no reference-agent
|
||||
analog**. On startup and before every trading tick the runner must answer one
|
||||
question before it is allowed to trade again: *does broker truth match the
|
||||
durable last-known state we persisted before the crash?* Trading is not
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Wall-clock scheduler for the live runtime (SPEC.md §7.5 #1).
|
||||
|
||||
Follows nanobot's ``cron/service.py`` shape: an asyncio timer that sleeps until
|
||||
Follows the reference ``cron/service.py`` shape: an asyncio timer that sleeps until
|
||||
the earliest ``next_run_at`` across all jobs, fires the due jobs, advances each
|
||||
fired job's ``next_run_at``, then re-sleeps. The sleep is **capped** at a
|
||||
re-check interval (default 5 min) so a far-future job still wakes the loop
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Live-runtime triggers — when an autonomous tick is due (SPEC §7.5 component 4).
|
||||
|
||||
This layers market-session, interval, and event triggers over R1's wall-clock
|
||||
scheduler. nanobot is time-only; trading additionally needs market-session
|
||||
scheduler. The reference scheduler is time-only; trading additionally needs market-session
|
||||
awareness ("is the US equity market open right now?") and event-driven
|
||||
predicates (price crossings, fill notifications). This module owns that layer.
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""OpenAI Codex OAuth provider.
|
||||
|
||||
This provider follows nanobot's OpenAI Codex OAuth path: a ChatGPT account is
|
||||
This provider follows the reference OpenAI Codex OAuth path: a ChatGPT account is
|
||||
authenticated by oauth-cli-kit, then requests are sent to the ChatGPT Codex
|
||||
Responses endpoint. It is intentionally separate from the standard OpenAI API
|
||||
key path because ChatGPT OAuth tokens are not OpenAI API keys.
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Security network compatibility helpers for channel adapters.
|
||||
|
||||
Re-exports validate_url_target / validate_resolved_url from :mod:`src.channels.utils`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.channels.utils import validate_resolved_url, validate_url_target
|
||||
|
||||
__all__ = ["validate_resolved_url", "validate_url_target"]
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Security workspace access compatibility helpers for channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
WORKSPACE_SCOPE_METADATA_KEY = "_workspace_scope"
|
||||
|
||||
|
||||
class WorkspaceScopeError(ValueError):
|
||||
"""Raised when a path is outside the allowed workspace scope."""
|
||||
|
||||
@property
|
||||
def message(self) -> str:
|
||||
"""Compatibility accessor for ported channel code."""
|
||||
return str(self)
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Security workspace policy compatibility helpers for channel adapters.
|
||||
|
||||
Re-exports :func:`is_path_within` from :mod:`src.channels.utils`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.channels.utils import is_path_within as _is_path_within
|
||||
|
||||
# Re-export for backward compat with ported channel code
|
||||
is_path_within = _is_path_within
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Goal-state payload helpers for WebSocket channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def goal_state_ws_blob(meta: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""Build a WebSocket payload for the current active goal state.
|
||||
|
||||
The WebSocket gateway stores goal snapshots in session metadata when they
|
||||
are available. This helper normalizes either a direct ``goal_state`` payload
|
||||
or compact goal metadata into the client event shape. It returns an inactive
|
||||
payload when no goal metadata is present.
|
||||
"""
|
||||
if not isinstance(meta, dict):
|
||||
return {"active": False, "active_goals": [], "completed_goals": []}
|
||||
goal_state = meta.get("goal_state")
|
||||
if isinstance(goal_state, dict):
|
||||
return {
|
||||
"active": bool(goal_state.get("active")),
|
||||
"active_goals": list(goal_state.get("active_goals") or []),
|
||||
"completed_goals": list(goal_state.get("completed_goals") or []),
|
||||
}
|
||||
goal = meta.get("active_goal") or meta.get("goal")
|
||||
if isinstance(goal, dict):
|
||||
return {"active": True, "active_goals": [goal], "completed_goals": []}
|
||||
return {"active": False, "active_goals": [], "completed_goals": []}
|
||||
@@ -0,0 +1,27 @@
|
||||
"""In-process WebUI turn timing helpers for channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
|
||||
_TURN_STARTED_AT: dict[str, float] = {}
|
||||
|
||||
|
||||
def mark_websocket_turn_started(chat_id: str, started_at: float | None = None) -> float:
|
||||
"""Record that a WebSocket chat turn started and return its timestamp."""
|
||||
value = float(started_at if started_at is not None else time.time())
|
||||
_TURN_STARTED_AT[str(chat_id)] = value
|
||||
return value
|
||||
|
||||
|
||||
def clear_websocket_turn_started(chat_id: str) -> None:
|
||||
"""Clear a recorded WebSocket turn start timestamp."""
|
||||
_TURN_STARTED_AT.pop(str(chat_id), None)
|
||||
|
||||
|
||||
def websocket_turn_wall_started_at(chat_id: str | None = None) -> float | None:
|
||||
"""Return the wall-clock time when the current turn started, or None."""
|
||||
if chat_id is None:
|
||||
return None
|
||||
return _TURN_STARTED_AT.get(str(chat_id))
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Media decoding helpers for channel adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class FileSizeExceeded(ValueError):
|
||||
"""Raised when a file exceeds the max size limit."""
|
||||
|
||||
_DATA_URL_RE = re.compile(r"^data:([^;,]+)((?:;[^,]*)*);base64,(.*)$", re.IGNORECASE | re.DOTALL)
|
||||
_EXT_BY_MIME = {
|
||||
"image/png": ".png",
|
||||
"image/jpeg": ".jpg",
|
||||
"image/jpg": ".jpg",
|
||||
"image/gif": ".gif",
|
||||
"image/webp": ".webp",
|
||||
"video/mp4": ".mp4",
|
||||
"video/webm": ".webm",
|
||||
"text/plain": ".txt",
|
||||
"application/pdf": ".pdf",
|
||||
}
|
||||
|
||||
|
||||
def save_base64_data_url(data_url: str, output_dir: Path, *, max_bytes: int = 0) -> Path:
|
||||
"""Decode a base64 data URL and save it to *output_dir*.
|
||||
|
||||
Args:
|
||||
data_url: A ``data:<mime>;base64,...`` URL.
|
||||
output_dir: Directory where the decoded file should be written.
|
||||
max_bytes: Optional maximum decoded byte length. ``0`` disables the
|
||||
limit.
|
||||
|
||||
Returns:
|
||||
The path of the decoded file.
|
||||
|
||||
Raises:
|
||||
FileSizeExceeded: If decoded data exceeds ``max_bytes``.
|
||||
ValueError: If the URL is malformed or uses an unsupported MIME type.
|
||||
"""
|
||||
match = _DATA_URL_RE.match(data_url)
|
||||
if not match:
|
||||
raise ValueError("expected data:<mime>;base64,<payload>")
|
||||
mime = match.group(1).strip().lower()
|
||||
ext = _EXT_BY_MIME.get(mime)
|
||||
if not ext:
|
||||
raise ValueError(f"unsupported data URL MIME type: {mime}")
|
||||
|
||||
payload = match.group(3).strip()
|
||||
try:
|
||||
decoded = base64.b64decode(payload, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError("invalid base64 payload") from exc
|
||||
|
||||
if max_bytes and len(decoded) > max_bytes:
|
||||
raise FileSizeExceeded(f"decoded media exceeds {max_bytes} bytes")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
path = output_dir / f"{uuid.uuid4().hex}{ext}"
|
||||
path.write_bytes(decoded)
|
||||
return path
|
||||
@@ -0,0 +1,76 @@
|
||||
"""API contracts for IM channel runtime controls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import api_server
|
||||
|
||||
|
||||
class FakeSessionService:
|
||||
"""SessionService placeholder for channel status tests."""
|
||||
|
||||
|
||||
def _client(tmp_path: Path, monkeypatch) -> TestClient:
|
||||
import src.channels.config as channel_config
|
||||
import src.channels.pairing.store as pairing_store
|
||||
|
||||
monkeypatch.setattr(api_server, "_channel_runtime", None)
|
||||
monkeypatch.setattr(api_server, "_channel_bus", None)
|
||||
monkeypatch.setattr(api_server, "_channel_manager", None)
|
||||
monkeypatch.setattr(api_server, "_get_session_service", lambda: FakeSessionService())
|
||||
monkeypatch.setattr(
|
||||
channel_config,
|
||||
"load_channels_config",
|
||||
lambda: {
|
||||
"send_max_retries": 1,
|
||||
"websocket": {"enabled": False, "allow_from": ["*"]},
|
||||
"telegram": {"enabled": False},
|
||||
"slack": {"enabled": True},
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(pairing_store, "_store_path", lambda: tmp_path / "pairing.json")
|
||||
return TestClient(api_server.app, client=("127.0.0.1", 50000))
|
||||
|
||||
|
||||
def test_channels_status_reports_all_configured_adapters(tmp_path: Path, monkeypatch) -> None:
|
||||
client = _client(tmp_path, monkeypatch)
|
||||
|
||||
response = client.get("/channels/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload["running"] is False
|
||||
assert payload["channels"]["websocket"]["configured"] is True
|
||||
assert payload["channels"]["telegram"]["enabled"] is False
|
||||
assert payload["channels"]["slack"]["enabled"] is True
|
||||
assert "available" in payload["channels"]["slack"]
|
||||
|
||||
|
||||
def test_channels_start_stop_are_controlled_runtime_actions(tmp_path: Path, monkeypatch) -> None:
|
||||
client = _client(tmp_path, monkeypatch)
|
||||
|
||||
start = client.post("/channels/start")
|
||||
stop = client.post("/channels/stop")
|
||||
|
||||
assert start.status_code == 200
|
||||
assert start.json()["status"] == "started"
|
||||
assert start.json()["running"] is True
|
||||
assert stop.status_code == 200
|
||||
assert stop.json()["status"] == "stopped"
|
||||
assert stop.json()["running"] is False
|
||||
|
||||
|
||||
def test_channels_pairing_command_uses_shared_store(tmp_path: Path, monkeypatch) -> None:
|
||||
client = _client(tmp_path, monkeypatch)
|
||||
|
||||
response = client.post(
|
||||
"/channels/pairing/command",
|
||||
json={"channel": "telegram", "command": "list"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["channel"] == "telegram"
|
||||
assert "No pending pairing requests" in response.json()["reply"]
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Contracts for the IM channel runtime wiring."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from src.channels.bus.events import InboundMessage, OutboundMessage
|
||||
from src.channels.bus.queue import MessageBus
|
||||
from src.channels.manager import ChannelManager
|
||||
from src.channels.registry import discover_channel_names, inspect_channels
|
||||
from src.channelsui.cli_apps_api import normalize_cli_app_mentions
|
||||
from src.channelsui.gateway_services import build_gateway_services
|
||||
from src.channelsui.mcp_presets_api import normalize_mcp_preset_mentions
|
||||
from src.channelsui.transcription_ws import webui_transcription_event
|
||||
from src.session.goal_state import goal_state_ws_blob
|
||||
from src.session.models import Message, Session
|
||||
from src.session.webui_turns import (
|
||||
clear_websocket_turn_started,
|
||||
mark_websocket_turn_started,
|
||||
websocket_turn_wall_started_at,
|
||||
)
|
||||
from src.utils.media_decode import FileSizeExceeded, save_base64_data_url
|
||||
|
||||
|
||||
class FakeSessionService:
|
||||
"""Small SessionService stand-in for channel runtime tests."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.created: list[Session] = []
|
||||
self.sent: list[tuple[str, str]] = []
|
||||
self.messages: dict[str, list[Message]] = {}
|
||||
|
||||
def create_session(self, title: str = "", config: dict[str, Any] | None = None) -> Session:
|
||||
session = Session(session_id=f"session-{len(self.created) + 1}", title=title, config=config or {})
|
||||
self.created.append(session)
|
||||
self.messages[session.session_id] = []
|
||||
return session
|
||||
|
||||
def get_session(self, session_id: str) -> Session | None:
|
||||
return next((session for session in self.created if session.session_id == session_id), None)
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
session_id: str,
|
||||
content: str,
|
||||
*,
|
||||
include_shell_tools: bool = False,
|
||||
parent_attempt_id: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
del include_shell_tools, parent_attempt_id
|
||||
self.sent.append((session_id, content))
|
||||
attempt_id = f"attempt-{len(self.sent)}"
|
||||
self.messages[session_id].append(
|
||||
Message(
|
||||
session_id=session_id,
|
||||
role="assistant",
|
||||
content=f"agent reply: {content}",
|
||||
linked_attempt_id=attempt_id,
|
||||
)
|
||||
)
|
||||
return {"message_id": "msg-1", "attempt_id": attempt_id}
|
||||
|
||||
def get_messages(self, session_id: str, limit: int = 100) -> list[Message]:
|
||||
del limit
|
||||
return list(self.messages.get(session_id, []))
|
||||
|
||||
|
||||
def test_channel_manager_can_construct_websocket_with_default_gateway() -> None:
|
||||
bus = MessageBus()
|
||||
manager = ChannelManager(
|
||||
{"websocket": {"enabled": True, "host": "127.0.0.1", "port": 0, "allow_from": ["*"]}},
|
||||
bus,
|
||||
)
|
||||
|
||||
assert "websocket" in manager.channels
|
||||
assert manager.channels["websocket"].name == "websocket"
|
||||
|
||||
|
||||
def test_registry_reports_all_built_in_channels_with_dependency_recovery() -> None:
|
||||
expected = {
|
||||
"dingtalk",
|
||||
"discord",
|
||||
"email",
|
||||
"feishu",
|
||||
"matrix",
|
||||
"mochat",
|
||||
"msteams",
|
||||
"napcat",
|
||||
"qq",
|
||||
"signal",
|
||||
"slack",
|
||||
"telegram",
|
||||
"websocket",
|
||||
"wecom",
|
||||
"weixin",
|
||||
"whatsapp",
|
||||
}
|
||||
|
||||
assert expected.issubset(set(discover_channel_names()))
|
||||
assert {"config", "runtime"}.isdisjoint(set(discover_channel_names()))
|
||||
|
||||
statuses = inspect_channels({"telegram": {"enabled": True}, "slack": {"enabled": True}})
|
||||
|
||||
for name in expected:
|
||||
assert name in statuses
|
||||
assert statuses[name]["name"] == name
|
||||
assert "install_hint" in statuses[name]
|
||||
assert "available" in statuses[name]
|
||||
|
||||
missing = [item for item in statuses.values() if not item["available"]]
|
||||
assert all(item["install_hint"].startswith("pip install") for item in missing)
|
||||
|
||||
|
||||
def test_channel_manager_status_includes_every_configured_adapter() -> None:
|
||||
bus = MessageBus()
|
||||
manager = ChannelManager(
|
||||
{
|
||||
"send_max_retries": 1,
|
||||
"websocket": {"enabled": True, "host": "127.0.0.1", "port": 0, "allow_from": ["*"]},
|
||||
"telegram": {"enabled": False},
|
||||
"slack": {"enabled": True},
|
||||
},
|
||||
bus,
|
||||
)
|
||||
|
||||
status = manager.get_status()
|
||||
|
||||
assert status["websocket"]["loaded"] is True
|
||||
assert status["websocket"]["enabled"] is True
|
||||
assert status["telegram"]["configured"] is True
|
||||
assert status["telegram"]["enabled"] is False
|
||||
assert status["slack"]["configured"] is True
|
||||
assert status["slack"]["enabled"] is True
|
||||
assert "available" in status["slack"]
|
||||
if not status["slack"]["loaded"]:
|
||||
assert status["slack"]["error"]
|
||||
assert status["slack"]["install_hint"].startswith("pip install")
|
||||
|
||||
|
||||
def test_registry_marks_lazy_sdk_adapter_unavailable(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import src.channels.discord as discord_channel
|
||||
|
||||
from src.channels.registry import inspect_channel
|
||||
|
||||
monkeypatch.setattr(discord_channel, "DISCORD_AVAILABLE", False)
|
||||
|
||||
status = inspect_channel("discord").to_dict()
|
||||
|
||||
assert status["available"] is False
|
||||
assert status["install_hint"] == "pip install 'vibe-trading-ai[discord]'"
|
||||
|
||||
|
||||
def test_channel_manager_skips_enabled_adapter_when_lazy_sdk_missing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import src.channels.discord as discord_channel
|
||||
|
||||
monkeypatch.setattr(discord_channel, "DISCORD_AVAILABLE", False)
|
||||
|
||||
manager = ChannelManager({"discord": {"enabled": True, "token": "x"}}, MessageBus())
|
||||
status = manager.get_status()["discord"]
|
||||
|
||||
assert "discord" not in manager.channels
|
||||
assert status["enabled"] is True
|
||||
assert status["available"] is False
|
||||
assert status["loaded"] is False
|
||||
assert status["install_hint"] == "pip install 'vibe-trading-ai[discord]'"
|
||||
|
||||
|
||||
def test_websocket_turn_wall_accepts_optional_chat_id() -> None:
|
||||
assert websocket_turn_wall_started_at("chat-1") is None
|
||||
mark_websocket_turn_started("chat-1", 123.0)
|
||||
assert websocket_turn_wall_started_at("chat-1") == 123.0
|
||||
clear_websocket_turn_started("chat-1")
|
||||
assert websocket_turn_wall_started_at("chat-1") is None
|
||||
|
||||
|
||||
def test_websocket_compatibility_helpers_have_structured_behavior() -> None:
|
||||
assert goal_state_ws_blob({"active_goal": {"goal_id": "g1"}}) == {
|
||||
"active": True,
|
||||
"active_goals": [{"goal_id": "g1"}],
|
||||
"completed_goals": [],
|
||||
}
|
||||
assert normalize_cli_app_mentions("@codex, terminal") == ["codex", "terminal"]
|
||||
assert normalize_mcp_preset_mentions(["@research", "market"]) == ["research", "market"]
|
||||
|
||||
|
||||
def test_gateway_services_adapt_session_service_for_hydration() -> None:
|
||||
service = FakeSessionService()
|
||||
session = service.create_session(title="chat", config={"metadata": {"active_goal": {"goal_id": "g1"}}})
|
||||
gateway = build_gateway_services(session_manager=service)
|
||||
|
||||
row = gateway.session_manager.read_session_file(session.session_id)
|
||||
|
||||
assert row["metadata"] == {"active_goal": {"goal_id": "g1"}}
|
||||
|
||||
|
||||
def test_transcription_event_returns_structured_error() -> None:
|
||||
async def scenario() -> None:
|
||||
event, payload = await webui_transcription_event({"type": "transcribe_audio"})
|
||||
assert event == "error"
|
||||
assert "not configured" in payload["detail"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_save_base64_data_url_decodes_and_limits_size(tmp_path: Path) -> None:
|
||||
payload = base64.b64encode(b"hello").decode("ascii")
|
||||
|
||||
saved = save_base64_data_url(f"data:text/plain;base64,{payload}", tmp_path, max_bytes=10)
|
||||
|
||||
assert saved.read_text(encoding="utf-8") == "hello"
|
||||
assert saved.suffix == ".txt"
|
||||
|
||||
with pytest.raises(FileSizeExceeded):
|
||||
save_base64_data_url(f"data:text/plain;base64,{payload}", tmp_path, max_bytes=4)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
save_base64_data_url("not-a-data-url", tmp_path, max_bytes=10)
|
||||
|
||||
|
||||
def test_channel_runtime_routes_inbound_to_session_and_outbound(tmp_path: Path) -> None:
|
||||
async def scenario() -> None:
|
||||
from src.channels.runtime import ChannelRuntime
|
||||
|
||||
bus = MessageBus()
|
||||
service = FakeSessionService()
|
||||
runtime = ChannelRuntime(
|
||||
bus=bus,
|
||||
session_service=service,
|
||||
manager=None,
|
||||
session_map_path=tmp_path / "channel_sessions.json",
|
||||
reply_timeout_s=1,
|
||||
poll_interval_s=0.01,
|
||||
)
|
||||
await runtime.start(start_manager=False)
|
||||
try:
|
||||
await bus.publish_inbound(
|
||||
InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user-1",
|
||||
chat_id="chat-1",
|
||||
content="hello from IM",
|
||||
)
|
||||
)
|
||||
|
||||
outbound = await asyncio.wait_for(bus.consume_outbound(), timeout=1)
|
||||
finally:
|
||||
await runtime.stop()
|
||||
|
||||
assert service.sent == [("session-1", "hello from IM")]
|
||||
assert outbound == OutboundMessage(
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
content="agent reply: hello from IM",
|
||||
metadata={
|
||||
"_channel_runtime": True,
|
||||
"attempt_id": "attempt-1",
|
||||
"session_id": "session-1",
|
||||
},
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_channel_runtime_handles_pairing_commands_without_agent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def scenario() -> None:
|
||||
from src.channels.runtime import ChannelRuntime
|
||||
|
||||
monkeypatch.setenv("VIBE_TRADING_DATA_DIR", str(tmp_path / "data"))
|
||||
bus = MessageBus()
|
||||
service = FakeSessionService()
|
||||
runtime = ChannelRuntime(
|
||||
bus=bus,
|
||||
session_service=service,
|
||||
manager=None,
|
||||
session_map_path=tmp_path / "channel_sessions.json",
|
||||
reply_timeout_s=1,
|
||||
poll_interval_s=0.01,
|
||||
)
|
||||
await runtime.start(start_manager=False)
|
||||
try:
|
||||
await bus.publish_inbound(
|
||||
InboundMessage(
|
||||
channel="telegram",
|
||||
sender_id="owner",
|
||||
chat_id="chat-1",
|
||||
content="/pairing list",
|
||||
)
|
||||
)
|
||||
outbound = await asyncio.wait_for(bus.consume_outbound(), timeout=1)
|
||||
finally:
|
||||
await runtime.stop()
|
||||
|
||||
assert service.sent == []
|
||||
assert outbound.channel == "telegram"
|
||||
assert outbound.chat_id == "chat-1"
|
||||
assert "No pending pairing requests" in outbound.content
|
||||
assert outbound.metadata["_pairing_command"] is True
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -0,0 +1,79 @@
|
||||
"""CLI coverage for IM channel commands."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from cli import _legacy
|
||||
|
||||
|
||||
def test_channels_parser_accepts_status_start_stop_login_and_pairing() -> None:
|
||||
parser = _legacy._build_parser()
|
||||
|
||||
status = parser.parse_args(["channels", "status", "--local", "--json"])
|
||||
assert status.command == "channels"
|
||||
assert status.channels_command == "status"
|
||||
assert status.local is True
|
||||
assert status.channels_json is True
|
||||
|
||||
start = parser.parse_args(["channels", "start"])
|
||||
assert start.channels_command == "start"
|
||||
|
||||
stop = parser.parse_args(["channels", "stop"])
|
||||
assert stop.channels_command == "stop"
|
||||
|
||||
login = parser.parse_args(["channels", "login", "weixin", "--force"])
|
||||
assert login.channels_command == "login"
|
||||
assert login.channel_name == "weixin"
|
||||
assert login.force is True
|
||||
|
||||
pairing = parser.parse_args(["channels", "pairing", "--channel", "telegram", "approve", "ABCD-EFGH"])
|
||||
assert pairing.channels_command == "pairing"
|
||||
assert pairing.channel == "telegram"
|
||||
assert pairing.pairing_command == "approve"
|
||||
assert pairing.pairing_args == ["ABCD-EFGH"]
|
||||
|
||||
|
||||
def test_channels_pairing_command_runs_against_local_store(tmp_path: Path, monkeypatch) -> None:
|
||||
import src.channels.pairing.store as pairing_store
|
||||
|
||||
monkeypatch.setattr(pairing_store, "_store_path", lambda: tmp_path / "pairing.json")
|
||||
|
||||
assert _legacy.main(["channels", "pairing", "--channel", "telegram", "list"]) == _legacy.EXIT_SUCCESS
|
||||
|
||||
|
||||
def test_channels_status_can_render_local_json(monkeypatch) -> None:
|
||||
import src.channels.config as channel_config
|
||||
|
||||
monkeypatch.setattr(
|
||||
channel_config,
|
||||
"load_channels_config",
|
||||
lambda: {"websocket": {"enabled": False}, "telegram": {"enabled": False}},
|
||||
)
|
||||
|
||||
assert _legacy.main(["channels", "status", "--local", "--json"]) == _legacy.EXIT_SUCCESS
|
||||
|
||||
|
||||
def test_channels_api_call_sends_configured_bearer_token(monkeypatch) -> None:
|
||||
captured = {}
|
||||
|
||||
class FakeResponse:
|
||||
def raise_for_status(self) -> None:
|
||||
return None
|
||||
|
||||
def json(self) -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
def fake_get(url, *, headers, timeout):
|
||||
captured["url"] = url
|
||||
captured["headers"] = headers
|
||||
captured["timeout"] = timeout
|
||||
return FakeResponse()
|
||||
|
||||
import httpx
|
||||
|
||||
monkeypatch.setenv("API_AUTH_KEY", "secret-token")
|
||||
monkeypatch.setattr(httpx, "get", fake_get)
|
||||
|
||||
assert _legacy._channels_api_call("GET", "/channels/status") == {"status": "ok"}
|
||||
assert captured["headers"] == {"Authorization": "Bearer secret-token"}
|
||||
@@ -44,3 +44,72 @@ def test_harmonic_backend_is_available_as_an_optional_extra() -> None:
|
||||
}
|
||||
|
||||
assert "pyharmonics" in harmonic_extra
|
||||
|
||||
|
||||
def test_channel_core_websocket_dependency_is_declared_for_baseline_installs() -> None:
|
||||
"""The built-in WebSocket gateway imports websockets at module import time."""
|
||||
pyproject = tomllib.loads((ROOT / "pyproject.toml").read_text())
|
||||
|
||||
core_dependencies = {
|
||||
_normalized_requirement_name(requirement)
|
||||
for requirement in pyproject["project"]["dependencies"]
|
||||
}
|
||||
requirements_txt = {
|
||||
_normalized_requirement_name(line)
|
||||
for line in (ROOT / "agent" / "requirements.txt").read_text().splitlines()
|
||||
if line and not line.startswith("#")
|
||||
}
|
||||
|
||||
assert "websockets" in core_dependencies
|
||||
assert "websockets" in requirements_txt
|
||||
|
||||
|
||||
def test_channel_optional_extras_cover_all_sdk_backed_adapters() -> None:
|
||||
"""Keep install hints and packaging extras in sync for IM adapters."""
|
||||
pyproject = tomllib.loads((ROOT / "pyproject.toml").read_text())
|
||||
extras = pyproject["project"]["optional-dependencies"]
|
||||
|
||||
expected_extras = {
|
||||
"channels",
|
||||
"dingtalk",
|
||||
"discord",
|
||||
"feishu",
|
||||
"matrix",
|
||||
"mochat",
|
||||
"msteams",
|
||||
"napcat",
|
||||
"qq",
|
||||
"slack",
|
||||
"telegram",
|
||||
"wecom",
|
||||
"weixin",
|
||||
"whatsapp",
|
||||
}
|
||||
assert expected_extras.issubset(set(extras))
|
||||
|
||||
channel_extra = {
|
||||
_normalized_requirement_name(requirement)
|
||||
for requirement in extras["channels"]
|
||||
}
|
||||
expected_packages = {
|
||||
"aiohttp",
|
||||
"cryptography",
|
||||
"dingtalk-stream",
|
||||
"discord.py",
|
||||
"lark-oapi",
|
||||
"matrix-nio",
|
||||
"mistune",
|
||||
"msgpack",
|
||||
"neonize",
|
||||
"nh3",
|
||||
"pyjwt",
|
||||
"python-socketio",
|
||||
"python-telegram-bot",
|
||||
"qrcode",
|
||||
"qq-botpy",
|
||||
"slack-sdk",
|
||||
"slackify-markdown",
|
||||
"websockets",
|
||||
"wecom-aibot-sdk",
|
||||
}
|
||||
assert expected_packages.issubset(channel_extra)
|
||||
@@ -253,6 +253,30 @@
|
||||
"useProviderDefaults": "استخدام الإعدادات الافتراضية للمزوّد",
|
||||
"modelIdHint": "استخدم معرّف النموذج الدقيق المطلوب من قِبل المزوّد.",
|
||||
"clearApiKey": "مسح مفتاح API المحفوظ",
|
||||
"channels": {
|
||||
"title": "قنوات IM",
|
||||
"description": "صل محولات الدردشة بنفس وقت تشغيل جلسات Vibe-Trading المستخدم في Web UI و CLI.",
|
||||
"runtime": "وقت التشغيل",
|
||||
"running": "قيد التشغيل",
|
||||
"stopped": "متوقف",
|
||||
"enabled": "مفعّل",
|
||||
"disabled": "معطّل",
|
||||
"loaded": "محمّل",
|
||||
"notLoaded": "غير محمّل",
|
||||
"unavailable": "غير متاح",
|
||||
"channel": "القناة",
|
||||
"state": "الحالة",
|
||||
"recovery": "الاسترداد",
|
||||
"noRecovery": "جاهز",
|
||||
"refresh": "تحديث",
|
||||
"start": "بدء القنوات",
|
||||
"stop": "إيقاف القنوات",
|
||||
"started": "تم بدء قنوات IM",
|
||||
"stoppedToast": "تم إيقاف قنوات IM",
|
||||
"refreshFailed": "فشل تحديث قنوات IM",
|
||||
"startFailed": "فشل بدء قنوات IM",
|
||||
"stopFailed": "فشل إيقاف قنوات IM"
|
||||
},
|
||||
"generation": "إعدادات التوليد",
|
||||
"dataSourceSettings": "إعدادات مصادر البيانات",
|
||||
"dataSourceSettingsDesc": "اضبط بيانات اعتماد بيانات السوق الاختيارية المستخدمة في الاختبارات العكسية ووكلاء البحث.",
|
||||
|
||||
@@ -258,6 +258,30 @@
|
||||
"useProviderDefaults": "Use provider defaults",
|
||||
"modelIdHint": "Use the exact model id required by your provider.",
|
||||
"clearApiKey": "Clear saved API key",
|
||||
"channels": {
|
||||
"title": "IM Channels",
|
||||
"description": "Connect chat adapters to the same Vibe-Trading session runtime used by the Web UI and CLI.",
|
||||
"runtime": "Runtime",
|
||||
"running": "Running",
|
||||
"stopped": "Stopped",
|
||||
"enabled": "Enabled",
|
||||
"disabled": "Disabled",
|
||||
"loaded": "Loaded",
|
||||
"notLoaded": "Not loaded",
|
||||
"unavailable": "Unavailable",
|
||||
"channel": "Channel",
|
||||
"state": "State",
|
||||
"recovery": "Recovery",
|
||||
"noRecovery": "Ready",
|
||||
"refresh": "Refresh",
|
||||
"start": "Start channels",
|
||||
"stop": "Stop channels",
|
||||
"started": "IM channels started",
|
||||
"stoppedToast": "IM channels stopped",
|
||||
"refreshFailed": "Failed to refresh IM channels",
|
||||
"startFailed": "Failed to start IM channels",
|
||||
"stopFailed": "Failed to stop IM channels"
|
||||
},
|
||||
"generation": "Generation",
|
||||
"dataSourceSettings": "Data Source Settings",
|
||||
"dataSourceSettingsDesc": "Configure optional market data credentials used by backtests and research agents.",
|
||||
|
||||
@@ -253,6 +253,30 @@
|
||||
"useProviderDefaults": "プロバイダーのデフォルトを使用",
|
||||
"modelIdHint": "プロバイダーが必要とする正確なモデル ID を使用してください。",
|
||||
"clearApiKey": "保存済み API キーをクリア",
|
||||
"channels": {
|
||||
"title": "IM チャンネル",
|
||||
"description": "チャットアダプターを Web UI と CLI が使う同じ Vibe-Trading セッションランタイムへ接続します。",
|
||||
"runtime": "ランタイム",
|
||||
"running": "実行中",
|
||||
"stopped": "停止中",
|
||||
"enabled": "有効",
|
||||
"disabled": "無効",
|
||||
"loaded": "ロード済み",
|
||||
"notLoaded": "未ロード",
|
||||
"unavailable": "利用不可",
|
||||
"channel": "チャンネル",
|
||||
"state": "状態",
|
||||
"recovery": "復旧",
|
||||
"noRecovery": "準備完了",
|
||||
"refresh": "更新",
|
||||
"start": "チャンネル開始",
|
||||
"stop": "チャンネル停止",
|
||||
"started": "IM チャンネルを開始しました",
|
||||
"stoppedToast": "IM チャンネルを停止しました",
|
||||
"refreshFailed": "IM チャンネルの更新に失敗しました",
|
||||
"startFailed": "IM チャンネルの開始に失敗しました",
|
||||
"stopFailed": "IM チャンネルの停止に失敗しました"
|
||||
},
|
||||
"generation": "生成設定",
|
||||
"dataSourceSettings": "データソース設定",
|
||||
"dataSourceSettingsDesc": "バックテストとリサーチエージェントが使用するオプションのマーケットデータ認証情報を設定します。",
|
||||
|
||||
@@ -253,6 +253,30 @@
|
||||
"useProviderDefaults": "프로바이더 기본값 사용",
|
||||
"modelIdHint": "프로바이더가 요구하는 정확한 모델 ID를 사용하세요.",
|
||||
"clearApiKey": "저장된 API 키 삭제",
|
||||
"channels": {
|
||||
"title": "IM 채널",
|
||||
"description": "채팅 어댑터를 Web UI와 CLI가 사용하는 동일한 Vibe-Trading 세션 런타임에 연결합니다.",
|
||||
"runtime": "런타임",
|
||||
"running": "실행 중",
|
||||
"stopped": "중지됨",
|
||||
"enabled": "활성화됨",
|
||||
"disabled": "비활성화됨",
|
||||
"loaded": "로드됨",
|
||||
"notLoaded": "로드되지 않음",
|
||||
"unavailable": "사용 불가",
|
||||
"channel": "채널",
|
||||
"state": "상태",
|
||||
"recovery": "복구",
|
||||
"noRecovery": "준비됨",
|
||||
"refresh": "새로 고침",
|
||||
"start": "채널 시작",
|
||||
"stop": "채널 중지",
|
||||
"started": "IM 채널이 시작되었습니다",
|
||||
"stoppedToast": "IM 채널이 중지되었습니다",
|
||||
"refreshFailed": "IM 채널 새로 고침 실패",
|
||||
"startFailed": "IM 채널 시작 실패",
|
||||
"stopFailed": "IM 채널 중지 실패"
|
||||
},
|
||||
"generation": "생성 설정",
|
||||
"dataSourceSettings": "데이터 소스 설정",
|
||||
"dataSourceSettingsDesc": "백테스트와 리서치 에이전트가 사용하는 선택적 마켓 데이터 자격 증명을 구성합니다.",
|
||||
|
||||
@@ -253,6 +253,30 @@
|
||||
"useProviderDefaults": "使用提供商默认设置",
|
||||
"modelIdHint": "使用提供商所需的准确模型 ID。",
|
||||
"clearApiKey": "清除已保存的 API 密钥",
|
||||
"channels": {
|
||||
"title": "IM 通道",
|
||||
"description": "将聊天适配器接入 Web UI 和 CLI 使用的同一套 Vibe-Trading 会话运行时。",
|
||||
"runtime": "运行时",
|
||||
"running": "运行中",
|
||||
"stopped": "已停止",
|
||||
"enabled": "已启用",
|
||||
"disabled": "已禁用",
|
||||
"loaded": "已加载",
|
||||
"notLoaded": "未加载",
|
||||
"unavailable": "不可用",
|
||||
"channel": "通道",
|
||||
"state": "状态",
|
||||
"recovery": "恢复方式",
|
||||
"noRecovery": "就绪",
|
||||
"refresh": "刷新",
|
||||
"start": "启动通道",
|
||||
"stop": "停止通道",
|
||||
"started": "IM 通道已启动",
|
||||
"stoppedToast": "IM 通道已停止",
|
||||
"refreshFailed": "刷新 IM 通道失败",
|
||||
"startFailed": "启动 IM 通道失败",
|
||||
"stopFailed": "停止 IM 通道失败"
|
||||
},
|
||||
"generation": "生成设置",
|
||||
"dataSourceSettings": "数据源设置",
|
||||
"dataSourceSettingsDesc": "配置回测引擎和研究智能体使用的可选市场数据凭证。",
|
||||
|
||||
@@ -143,6 +143,14 @@ export const api = {
|
||||
method: "PUT",
|
||||
body: JSON.stringify(settings),
|
||||
}),
|
||||
getChannelStatus: () => request<ChannelRuntimeStatus>("/channels/status"),
|
||||
startChannels: () => request<ChannelRuntimeActionResponse>("/channels/start", { method: "POST" }),
|
||||
stopChannels: () => request<ChannelRuntimeActionResponse>("/channels/stop", { method: "POST" }),
|
||||
runChannelPairingCommand: (body: ChannelPairingCommandRequest) =>
|
||||
request<ChannelPairingCommandResponse>("/channels/pairing/command", {
|
||||
method: "POST",
|
||||
body: JSON.stringify(body),
|
||||
}),
|
||||
|
||||
// Alpha Zoo API
|
||||
listAlphas: (params: AlphaListParams = {}) => {
|
||||
@@ -278,6 +286,40 @@ export interface UpdateDataSourceSettingsRequest {
|
||||
clear_tushare_token?: boolean;
|
||||
}
|
||||
|
||||
export interface ChannelAdapterStatus {
|
||||
name: string;
|
||||
display_name: string;
|
||||
configured: boolean;
|
||||
enabled: boolean;
|
||||
available: boolean;
|
||||
loaded: boolean;
|
||||
running: boolean;
|
||||
error?: string;
|
||||
install_hint?: string;
|
||||
}
|
||||
|
||||
export interface ChannelRuntimeStatus {
|
||||
running: boolean;
|
||||
inbound_queue: number;
|
||||
outbound_queue: number;
|
||||
session_count: number;
|
||||
channels: Record<string, ChannelAdapterStatus>;
|
||||
}
|
||||
|
||||
export interface ChannelRuntimeActionResponse extends ChannelRuntimeStatus {
|
||||
status: string;
|
||||
}
|
||||
|
||||
export interface ChannelPairingCommandRequest {
|
||||
channel: string;
|
||||
command: string;
|
||||
}
|
||||
|
||||
export interface ChannelPairingCommandResponse {
|
||||
channel: string;
|
||||
reply: string;
|
||||
}
|
||||
|
||||
// --- Types matching backend API contracts ---
|
||||
|
||||
export interface RunListItem {
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
import i18n from "@/i18n";
|
||||
import { useEffect, useMemo, useState, type FormEvent } from "react";
|
||||
import { Database, KeyRound, Loader2, RotateCcw, Save, Server, SlidersHorizontal } from "lucide-react";
|
||||
import { Database, KeyRound, Loader2, MessageSquareMore, Play, RefreshCw, RotateCcw, Save, Server, SlidersHorizontal, Square } from "lucide-react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { toast } from "sonner";
|
||||
import { api, isAuthRequiredError, type DataSourceSettings, type LLMProviderOption, type LLMSettings } from "@/lib/api";
|
||||
import { api, isAuthRequiredError, type ChannelRuntimeStatus, type DataSourceSettings, type LLMProviderOption, type LLMSettings } from "@/lib/api";
|
||||
import { getApiAuthKey, setApiAuthKey } from "@/lib/apiAuth";
|
||||
|
||||
interface LLMFormState {
|
||||
@@ -33,9 +34,10 @@ function toForm(settings: LLMSettings): LLMFormState {
|
||||
}
|
||||
|
||||
export function Settings() {
|
||||
|
||||
const { t } = useTranslation();
|
||||
const [settings, setSettings] = useState<LLMSettings | null>(null);
|
||||
const [dataSettings, setDataSettings] = useState<DataSourceSettings | null>(null);
|
||||
const [channelStatus, setChannelStatus] = useState<ChannelRuntimeStatus | null>(null);
|
||||
const [form, setForm] = useState<LLMFormState | null>(null);
|
||||
const [apiKey, setApiKey] = useState("");
|
||||
const [localApiKey, setLocalApiKeyState] = useState(() => getApiAuthKey());
|
||||
@@ -45,16 +47,19 @@ export function Settings() {
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const [dataSaving, setDataSaving] = useState(false);
|
||||
const [channelRefreshing, setChannelRefreshing] = useState(false);
|
||||
const [channelAction, setChannelAction] = useState<"start" | "stop" | null>(null);
|
||||
const [settingsLoadError, setSettingsLoadError] = useState<string | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
let alive = true;
|
||||
Promise.all([api.getLLMSettings(), api.getDataSourceSettings()])
|
||||
.then(([llmData, dataSourceData]) => {
|
||||
Promise.all([api.getLLMSettings(), api.getDataSourceSettings(), api.getChannelStatus()])
|
||||
.then(([llmData, dataSourceData, channelData]) => {
|
||||
if (!alive) return;
|
||||
setSettings(llmData);
|
||||
setForm(toForm(llmData));
|
||||
setDataSettings(dataSourceData);
|
||||
setChannelStatus(channelData);
|
||||
setSettingsLoadError(null);
|
||||
})
|
||||
.catch((error) => {
|
||||
@@ -73,6 +78,30 @@ export function Settings() {
|
||||
return () => { alive = false; };
|
||||
}, []);
|
||||
|
||||
const refreshChannelStatus = async () => {
|
||||
setChannelRefreshing(true);
|
||||
try {
|
||||
setChannelStatus(await api.getChannelStatus());
|
||||
} catch (error) {
|
||||
toast.error(`${t("settings.channels.refreshFailed")}: ${error instanceof Error ? error.message : "Unknown error"}`);
|
||||
} finally {
|
||||
setChannelRefreshing(false);
|
||||
}
|
||||
};
|
||||
|
||||
const setChannelsRunning = async (action: "start" | "stop") => {
|
||||
setChannelAction(action);
|
||||
try {
|
||||
const updated = action === "start" ? await api.startChannels() : await api.stopChannels();
|
||||
setChannelStatus(updated);
|
||||
toast.success(action === "start" ? t("settings.channels.started") : t("settings.channels.stoppedToast"));
|
||||
} catch (error) {
|
||||
toast.error(`${action === "start" ? t("settings.channels.startFailed") : t("settings.channels.stopFailed")}: ${error instanceof Error ? error.message : "Unknown error"}`);
|
||||
} finally {
|
||||
setChannelAction(null);
|
||||
}
|
||||
};
|
||||
|
||||
const providers = settings?.providers ?? [];
|
||||
const selectedProvider = useMemo<LLMProviderOption | undefined>(
|
||||
() => providers.find((provider) => provider.name === form?.provider),
|
||||
@@ -182,7 +211,7 @@ export function Settings() {
|
||||
</form>
|
||||
);
|
||||
|
||||
if (loading || !form || !settings || !dataSettings) {
|
||||
if (loading || !form || !settings || !dataSettings || !channelStatus) {
|
||||
return (
|
||||
<div className="mx-auto max-w-5xl space-y-6 p-6">
|
||||
<div className="space-y-2">
|
||||
@@ -218,6 +247,11 @@ export function Settings() {
|
||||
const tushareStatus = dataSettings.tushare_token_configured
|
||||
? "Configured"
|
||||
: "Leave blank to keep the current token";
|
||||
const channelRows = Object.entries(channelStatus.channels ?? {}).sort(([a], [b]) => a.localeCompare(b));
|
||||
const channelEnabledCount = channelRows.filter(([, item]) => item.enabled).length;
|
||||
const channelLoadedCount = channelRows.filter(([, item]) => item.loaded).length;
|
||||
const channelUnavailableCount = channelRows.filter(([, item]) => item.available === false).length;
|
||||
const channelBusy = channelRefreshing || channelAction !== null;
|
||||
|
||||
return (
|
||||
<div className="mx-auto max-w-5xl space-y-6 p-6">
|
||||
@@ -228,6 +262,104 @@ export function Settings() {
|
||||
|
||||
{localApiAccessSection}
|
||||
|
||||
<section className="rounded-lg border bg-card p-5 shadow-sm">
|
||||
<div className="mb-5 flex flex-col gap-4 md:flex-row md:items-start md:justify-between">
|
||||
<div className="space-y-1">
|
||||
<div className="flex items-center gap-2">
|
||||
<MessageSquareMore className="h-4 w-4 text-primary" />
|
||||
<h2 className="text-base font-semibold">{t("settings.channels.title")}</h2>
|
||||
</div>
|
||||
<p className="max-w-3xl text-sm text-muted-foreground">{t("settings.channels.description")}</p>
|
||||
</div>
|
||||
<div className="flex flex-wrap gap-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={refreshChannelStatus}
|
||||
disabled={channelBusy}
|
||||
className="inline-flex items-center justify-center gap-2 rounded-md border px-3 py-2 text-sm text-muted-foreground transition hover:bg-muted hover:text-foreground disabled:cursor-not-allowed disabled:opacity-60"
|
||||
>
|
||||
{channelRefreshing ? <Loader2 className="h-4 w-4 animate-spin" /> : <RefreshCw className="h-4 w-4" />}
|
||||
{t("settings.channels.refresh")}
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setChannelsRunning("start")}
|
||||
disabled={channelBusy}
|
||||
className="inline-flex items-center justify-center gap-2 rounded-md bg-primary px-3 py-2 text-sm font-medium text-primary-foreground transition hover:opacity-90 disabled:cursor-not-allowed disabled:opacity-60"
|
||||
>
|
||||
{channelAction === "start" ? <Loader2 className="h-4 w-4 animate-spin" /> : <Play className="h-4 w-4" />}
|
||||
{t("settings.channels.start")}
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setChannelsRunning("stop")}
|
||||
disabled={channelBusy}
|
||||
className="inline-flex items-center justify-center gap-2 rounded-md border px-3 py-2 text-sm text-muted-foreground transition hover:bg-muted hover:text-foreground disabled:cursor-not-allowed disabled:opacity-60"
|
||||
>
|
||||
{channelAction === "stop" ? <Loader2 className="h-4 w-4 animate-spin" /> : <Square className="h-4 w-4" />}
|
||||
{t("settings.channels.stop")}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="mb-4 grid gap-3 md:grid-cols-4">
|
||||
<div className="rounded-md border bg-muted/20 px-3 py-2">
|
||||
<div className="text-xs text-muted-foreground">{t("settings.channels.runtime")}</div>
|
||||
<div className="text-sm font-medium">{channelStatus.running ? t("settings.channels.running") : t("settings.channels.stopped")}</div>
|
||||
</div>
|
||||
<div className="rounded-md border bg-muted/20 px-3 py-2">
|
||||
<div className="text-xs text-muted-foreground">{t("settings.channels.enabled")}</div>
|
||||
<div className="text-sm font-medium">{channelEnabledCount}</div>
|
||||
</div>
|
||||
<div className="rounded-md border bg-muted/20 px-3 py-2">
|
||||
<div className="text-xs text-muted-foreground">{t("settings.channels.loaded")}</div>
|
||||
<div className="text-sm font-medium">{channelLoadedCount}</div>
|
||||
</div>
|
||||
<div className="rounded-md border bg-muted/20 px-3 py-2">
|
||||
<div className="text-xs text-muted-foreground">{t("settings.channels.unavailable")}</div>
|
||||
<div className="text-sm font-medium">{channelUnavailableCount}</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="overflow-hidden rounded-md border">
|
||||
<table className="w-full text-sm">
|
||||
<thead className="bg-muted/40 text-xs text-muted-foreground">
|
||||
<tr>
|
||||
<th className="px-3 py-2 text-left font-medium">{t("settings.channels.channel")}</th>
|
||||
<th className="px-3 py-2 text-left font-medium">{t("settings.channels.state")}</th>
|
||||
<th className="px-3 py-2 text-left font-medium">{t("settings.channels.recovery")}</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{channelRows.map(([name, item]) => (
|
||||
<tr key={name} className="border-t">
|
||||
<td className="px-3 py-2 align-top">
|
||||
<div className="font-medium">{item.display_name || name}</div>
|
||||
<div className="text-xs text-muted-foreground">{name}</div>
|
||||
</td>
|
||||
<td className="px-3 py-2 align-top">
|
||||
<div className="flex flex-wrap gap-1.5">
|
||||
<span className={`rounded-full px-2 py-0.5 text-xs ${item.enabled ? "bg-primary/10 text-primary" : "bg-muted text-muted-foreground"}`}>
|
||||
{item.enabled ? t("settings.channels.enabled") : t("settings.channels.disabled")}
|
||||
</span>
|
||||
<span className={`rounded-full px-2 py-0.5 text-xs ${item.loaded ? "bg-success/10 text-success" : "bg-muted text-muted-foreground"}`}>
|
||||
{item.loaded ? t("settings.channels.loaded") : t("settings.channels.notLoaded")}
|
||||
</span>
|
||||
<span className={`rounded-full px-2 py-0.5 text-xs ${item.running ? "bg-success/10 text-success" : "bg-muted text-muted-foreground"}`}>
|
||||
{item.running ? t("settings.channels.running") : t("settings.channels.stopped")}
|
||||
</span>
|
||||
</div>
|
||||
</td>
|
||||
<td className="max-w-md px-3 py-2 align-top text-xs text-muted-foreground">
|
||||
{item.install_hint || item.error || t("settings.channels.noRecovery")}
|
||||
</td>
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<div className="space-y-2">
|
||||
<h2 className="text-lg font-semibold tracking-tight">{"LLM Settings"}</h2>
|
||||
<p className="max-w-3xl text-sm text-muted-foreground">{"Choose the model used by the agent and save it to the project-local agent/.env file."}</p>
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import { Settings } from "../Settings";
|
||||
|
||||
const apiMock = vi.hoisted(() => ({
|
||||
getLLMSettings: vi.fn(),
|
||||
getDataSourceSettings: vi.fn(),
|
||||
getChannelStatus: vi.fn(),
|
||||
startChannels: vi.fn(),
|
||||
stopChannels: vi.fn(),
|
||||
updateLLMSettings: vi.fn(),
|
||||
updateDataSourceSettings: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/lib/api", async () => {
|
||||
const actual = await vi.importActual<typeof import("@/lib/api")>("@/lib/api");
|
||||
return {
|
||||
...actual,
|
||||
api: apiMock,
|
||||
isAuthRequiredError: vi.fn(() => false),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("@/lib/apiAuth", () => ({
|
||||
getApiAuthKey: vi.fn(() => ""),
|
||||
setApiAuthKey: vi.fn(),
|
||||
}));
|
||||
|
||||
function llmSettings() {
|
||||
return {
|
||||
provider: "openrouter",
|
||||
model_name: "deepseek/deepseek-v3.2",
|
||||
base_url: "https://openrouter.ai/api/v1",
|
||||
api_key_env: "OPENROUTER_API_KEY",
|
||||
api_key_configured: false,
|
||||
api_key_required: true,
|
||||
temperature: 0.1,
|
||||
timeout_seconds: 120,
|
||||
max_retries: 2,
|
||||
reasoning_effort: "",
|
||||
sse_timeout_seconds: 300,
|
||||
env_path: "agent/.env",
|
||||
providers: [
|
||||
{
|
||||
name: "openrouter",
|
||||
label: "OpenRouter",
|
||||
api_key_env: "OPENROUTER_API_KEY",
|
||||
base_url_env: "OPENROUTER_BASE_URL",
|
||||
default_model: "deepseek/deepseek-v3.2",
|
||||
default_base_url: "https://openrouter.ai/api/v1",
|
||||
api_key_required: true,
|
||||
auth_type: "api_key",
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
function dataSourceSettings() {
|
||||
return {
|
||||
tushare_token_configured: false,
|
||||
baostock_supported: true,
|
||||
baostock_installed: true,
|
||||
baostock_message: "BaoStock available",
|
||||
env_path: "agent/.env",
|
||||
};
|
||||
}
|
||||
|
||||
function channelStatus(overrides = {}) {
|
||||
return {
|
||||
running: false,
|
||||
inbound_queue: 0,
|
||||
outbound_queue: 0,
|
||||
session_count: 0,
|
||||
channels: {
|
||||
websocket: {
|
||||
name: "websocket",
|
||||
display_name: "WebSocket",
|
||||
configured: true,
|
||||
enabled: true,
|
||||
available: true,
|
||||
loaded: true,
|
||||
running: false,
|
||||
error: "",
|
||||
install_hint: "",
|
||||
},
|
||||
telegram: {
|
||||
name: "telegram",
|
||||
display_name: "Telegram",
|
||||
configured: true,
|
||||
enabled: false,
|
||||
available: false,
|
||||
loaded: false,
|
||||
running: false,
|
||||
error: "ModuleNotFoundError",
|
||||
install_hint: "pip install 'vibe-trading-ai[telegram]'",
|
||||
},
|
||||
},
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("Settings IM channels panel", () => {
|
||||
beforeEach(() => {
|
||||
apiMock.getLLMSettings.mockResolvedValue(llmSettings());
|
||||
apiMock.getDataSourceSettings.mockResolvedValue(dataSourceSettings());
|
||||
apiMock.getChannelStatus.mockResolvedValue(channelStatus());
|
||||
apiMock.startChannels.mockResolvedValue(channelStatus({ running: true }));
|
||||
apiMock.stopChannels.mockResolvedValue(channelStatus());
|
||||
});
|
||||
|
||||
it("renders channel runtime status and refreshes it", async () => {
|
||||
render(<Settings />);
|
||||
|
||||
expect(await screen.findByText("IM Channels")).toBeInTheDocument();
|
||||
expect(screen.getByText("websocket")).toBeInTheDocument();
|
||||
expect(screen.getByText("telegram")).toBeInTheDocument();
|
||||
expect(screen.getByText("pip install 'vibe-trading-ai[telegram]'")).toBeInTheDocument();
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Refresh" }));
|
||||
|
||||
await waitFor(() => expect(apiMock.getChannelStatus).toHaveBeenCalledTimes(2));
|
||||
});
|
||||
|
||||
it("starts channels from the settings control surface", async () => {
|
||||
render(<Settings />);
|
||||
await screen.findByText("IM Channels");
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Start channels" }));
|
||||
|
||||
await waitFor(() => expect(apiMock.startChannels).toHaveBeenCalledTimes(1));
|
||||
});
|
||||
});
|
||||
@@ -47,6 +47,7 @@ dependencies = [
|
||||
"ccxt>=4.0.0",
|
||||
"fastapi>=0.104.0",
|
||||
"uvicorn[standard]>=0.24.0",
|
||||
"websockets>=12.0",
|
||||
"pydantic>=2.0.0",
|
||||
"python-multipart>=0.0.18",
|
||||
"sse-starlette>=1.6.0",
|
||||
@@ -107,6 +108,75 @@ ashare = [
|
||||
harmonic = [
|
||||
"pyharmonics>=1.5.0",
|
||||
]
|
||||
dingtalk = [
|
||||
"dingtalk-stream>=0.24.0",
|
||||
]
|
||||
discord = [
|
||||
"discord.py>=2.4.0",
|
||||
]
|
||||
feishu = [
|
||||
"lark-oapi>=1.4.0",
|
||||
"qrcode>=7.4.2",
|
||||
]
|
||||
matrix = [
|
||||
"aiohttp>=3.9.0",
|
||||
"matrix-nio[e2e]>=0.25.0",
|
||||
"mistune>=3.0.0",
|
||||
"nh3>=0.2.18",
|
||||
]
|
||||
mochat = [
|
||||
"msgpack>=1.0.0",
|
||||
"python-socketio[asyncio_client]>=5.11.0",
|
||||
]
|
||||
msteams = [
|
||||
"cryptography>=42.0.0",
|
||||
"PyJWT[crypto]>=2.8.0",
|
||||
]
|
||||
napcat = [
|
||||
"aiohttp>=3.9.0",
|
||||
"websockets>=12.0",
|
||||
]
|
||||
qq = [
|
||||
"aiohttp>=3.9.0",
|
||||
"qq-botpy>=1.2.0",
|
||||
]
|
||||
slack = [
|
||||
"slack-sdk>=3.33.0",
|
||||
"slackify-markdown>=4.4.0",
|
||||
]
|
||||
telegram = [
|
||||
"python-telegram-bot>=21.0",
|
||||
]
|
||||
wecom = [
|
||||
"wecom-aibot-sdk>=1.0.0",
|
||||
]
|
||||
weixin = [
|
||||
"qrcode>=7.4.2",
|
||||
]
|
||||
whatsapp = [
|
||||
"neonize>=0.3.18.post0",
|
||||
]
|
||||
channels = [
|
||||
"aiohttp>=3.9.0",
|
||||
"cryptography>=42.0.0",
|
||||
"dingtalk-stream>=0.24.0",
|
||||
"discord.py>=2.4.0",
|
||||
"lark-oapi>=1.4.0",
|
||||
"matrix-nio[e2e]>=0.25.0",
|
||||
"mistune>=3.0.0",
|
||||
"msgpack>=1.0.0",
|
||||
"neonize>=0.3.18.post0",
|
||||
"nh3>=0.2.18",
|
||||
"PyJWT[crypto]>=2.8.0",
|
||||
"python-socketio[asyncio_client]>=5.11.0",
|
||||
"python-telegram-bot>=21.0",
|
||||
"qrcode>=7.4.2",
|
||||
"qq-botpy>=1.2.0",
|
||||
"slack-sdk>=3.33.0",
|
||||
"slackify-markdown>=4.4.0",
|
||||
"websockets>=12.0",
|
||||
"wecom-aibot-sdk>=1.0.0",
|
||||
]
|
||||
dev = [
|
||||
"pytest>=7.0",
|
||||
"pytest-cov>=4.0",
|
||||
|
||||
Reference in new issue
Block a user