feat: wire IM channel runtime

Merge IM channel runtime wiring, controls, dependency gates, tests, and legacy-name cleanup.
This commit is contained in:
Haozhe Wu authored and GitHub committed 2026-06-30 14:31:26 +08:00
1 parent 6f6b1794a3
commit 101306e916
73 files changed
+20558 -27

No files matched your search

+3 -1
View File
@@ -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
+2 -2
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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>
+85
View File
@@ -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
+198
View File
@@ -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 -2
View File
@@ -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
View File
@@ -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
View File
@@ -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).
"""
+1
View File
@@ -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
+48
View File
@@ -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",
]
+237
View File
@@ -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
+6
View File
@@ -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"]
+54
View File
@@ -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)
+43
View File
@@ -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()
+21
View File
@@ -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)
+773
View File
@@ -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
+818
View File
@@ -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
+914
View File
@@ -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
+478
View File
@@ -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
+943
View File
@@ -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
+822
View File
@@ -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
+579
View File
@@ -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)
+33
View File
@@ -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",
]
+262
View File
@@ -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>]`"
)
+701
View File
@@ -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)
+262
View File
@@ -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)
+215
View File
@@ -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
+740
View File
@@ -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("&amp;", "&"), 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
+156
View File
@@ -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
+554
View File
@@ -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
+682
View File
@@ -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()
+1
View File
@@ -0,0 +1 @@
"""WebUI compatibility helpers for channel adapters."""
+19
View File
@@ -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()]
+23
View File
@@ -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
+301
View File
@@ -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
+78
View File
@@ -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")
+19
View File
@@ -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()]
+11
View File
@@ -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")
+12 -1
View File
@@ -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
+17
View File
@@ -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 -1
View File
@@ -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 -1
View File
@@ -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 -1
View File
@@ -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 -1
View File
@@ -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.
+10
View File
@@ -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"]
+14
View File
@@ -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)
+11
View File
@@ -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
+28
View File
@@ -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": []}
+27
View File
@@ -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))
+65
View File
@@ -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
+76
View File
@@ -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"]
+304
View File
@@ -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())
+79
View File
@@ -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)
+24
View File
@@ -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": "اضبط بيانات اعتماد بيانات السوق الاختيارية المستخدمة في الاختبارات العكسية ووكلاء البحث.",
+24
View File
@@ -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.",
+24
View File
@@ -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": "バックテストとリサーチエージェントが使用するオプションのマーケットデータ認証情報を設定します。",
+24
View File
@@ -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": "백테스트와 리서치 에이전트가 사용하는 선택적 마켓 데이터 자격 증명을 구성합니다.",
+24
View File
@@ -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": "配置回测引擎和研究智能体使用的可选市场数据凭证。",
+42
View File
@@ -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 {
+138 -6
View File
@@ -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));
});
});
+70
View File
@@ -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",