| 1 | """Channel manager for coordinating chat channels.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import asyncio |
| 6 | from pathlib import Path |
| 7 | from typing import TYPE_CHECKING, Any |
| 8 | |
| 9 | from loguru import logger |
| 10 | |
| 11 | from nanobot.bus.events import OutboundMessage |
| 12 | from nanobot.bus.queue import MessageBus |
| 13 | from nanobot.channels.base import BaseChannel |
| 14 | from nanobot.config.schema import Config |
| 15 | from nanobot.utils.restart import consume_restart_notice_from_env, format_restart_completed_message |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from nanobot.providers.base import LLMProvider |
| 19 | from nanobot.session.manager import SessionManager |
| 20 | |
| 21 | |
| 22 | def _default_webui_dist() -> Path | None: |
| 23 | """Return the absolute path to the bundled webui dist directory if it exists.""" |
| 24 | try: |
| 25 | import nanobot.web as web_pkg # type: ignore[import-not-found] |
| 26 | except ImportError: |
| 27 | return None |
| 28 | candidate = Path(web_pkg.__file__).resolve().parent / "dist" |
| 29 | return candidate if candidate.is_dir() else None |
| 30 | |
| 31 | # Retry delays for message sending (exponential backoff: 1s, 2s, 4s) |
| 32 | _SEND_RETRY_DELAYS = (1, 2, 4) |
| 33 | |
| 34 | |
| 35 | class ChannelManager: |
| 36 | """ |
| 37 | Manages chat channels and coordinates message routing. |
| 38 | |
| 39 | Responsibilities: |
| 40 | - Initialize enabled channels (Telegram, WhatsApp, etc.) |
| 41 | - Start/stop channels |
| 42 | - Route outbound messages |
| 43 | """ |
| 44 | |
| 45 | def __init__( |
| 46 | self, |
| 47 | config: Config, |
| 48 | bus: MessageBus, |
| 49 | *, |
| 50 | session_manager: "SessionManager | None" = None, |
| 51 | provider: "LLMProvider | None" = None, |
| 52 | gateway_debug: bool = False, |
| 53 | ): |
| 54 | self.config = config |
| 55 | self.bus = bus |
| 56 | self._session_manager = session_manager |
| 57 | self._provider = provider |
| 58 | self._gateway_debug = gateway_debug |
| 59 | self.channels: dict[str, BaseChannel] = {} |
| 60 | self._dispatch_task: asyncio.Task | None = None |
| 61 | |
| 62 | self._init_channels() |
| 63 | |
| 64 | def _configured_memory_review_runner(self): |
| 65 | """Bind Memory VLM calls to the same configured model gateway.""" |
| 66 | from functools import partial |
| 67 | |
| 68 | from nanobot.director.memory_coordinator import ( |
| 69 | run_memory_review_from_config, |
| 70 | ) |
| 71 | |
| 72 | review = self.config.tools.memory_review |
| 73 | model = review.model or self.config.agents.defaults.model |
| 74 | provider = review.provider or None |
| 75 | return partial( |
| 76 | run_memory_review_from_config, |
| 77 | api_base=self.config.get_api_base(model, provider=provider), |
| 78 | api_key=self.config.get_api_key(model, provider=provider) or "", |
| 79 | vlm_model=model, |
| 80 | candidate_count=review.candidate_count, |
| 81 | ) |
| 82 | |
| 83 | def _init_channels(self) -> None: |
| 84 | """Initialize channels discovered via pkgutil scan + entry_points plugins.""" |
| 85 | from nanobot.channels.registry import discover_all |
| 86 | |
| 87 | transcription_provider = self.config.channels.transcription_provider |
| 88 | transcription_key = self._resolve_transcription_key(transcription_provider) |
| 89 | transcription_base = self._resolve_transcription_base(transcription_provider) |
| 90 | transcription_language = self.config.channels.transcription_language |
| 91 | |
| 92 | for name, cls in discover_all().items(): |
| 93 | section = getattr(self.config.channels, name, None) |
| 94 | if section is None: |
| 95 | continue |
| 96 | enabled = ( |
| 97 | section.get("enabled", False) |
| 98 | if isinstance(section, dict) |
| 99 | else getattr(section, "enabled", False) |
| 100 | ) |
| 101 | if not enabled: |
| 102 | continue |
| 103 | try: |
| 104 | kwargs: dict[str, Any] = {} |
| 105 | # Only the WebSocket channel currently hosts the embedded webui |
| 106 | # surface; other channels stay oblivious to these knobs. |
| 107 | if cls.name == "websocket" and self._session_manager is not None: |
| 108 | kwargs["session_manager"] = self._session_manager |
| 109 | kwargs["provider"] = self._provider |
| 110 | kwargs["model"] = self.config.agents.defaults.model |
| 111 | kwargs["tools_config"] = self.config.tools |
| 112 | kwargs["gateway_debug"] = self._gateway_debug |
| 113 | kwargs["memory_review_runner"] = ( |
| 114 | self._configured_memory_review_runner() |
| 115 | ) |
| 116 | static_path = _default_webui_dist() |
| 117 | if static_path is not None: |
| 118 | kwargs["static_dist_path"] = static_path |
| 119 | if cls.name == "director_callback" and self._session_manager is not None: |
| 120 | kwargs["workspace"] = self._session_manager.workspace |
| 121 | kwargs["tools_config"] = self.config.tools |
| 122 | kwargs["memory_review_runner"] = ( |
| 123 | self._configured_memory_review_runner() |
| 124 | ) |
| 125 | channel = cls(section, self.bus, **kwargs) |
| 126 | channel.transcription_provider = transcription_provider |
| 127 | channel.transcription_api_key = transcription_key |
| 128 | channel.transcription_api_base = transcription_base |
| 129 | channel.transcription_language = transcription_language |
| 130 | self.channels[name] = channel |
| 131 | logger.info("{} channel enabled", cls.display_name) |
| 132 | except Exception as e: |
| 133 | logger.warning("{} channel not available: {}", name, e) |
| 134 | |
| 135 | self._validate_allow_from() |
| 136 | |
| 137 | def _resolve_transcription_key(self, provider: str) -> str: |
| 138 | """Pick the API key for the configured transcription provider.""" |
| 139 | try: |
| 140 | if provider == "openai": |
| 141 | return self.config.providers.openai.api_key |
| 142 | return self.config.providers.groq.api_key |
| 143 | except AttributeError: |
| 144 | return "" |
| 145 | |
| 146 | def _resolve_transcription_base(self, provider: str) -> str: |
| 147 | """Pick the API base URL for the configured transcription provider.""" |
| 148 | try: |
| 149 | if provider == "openai": |
| 150 | return self.config.providers.openai.api_base or "" |
| 151 | return self.config.providers.groq.api_base or "" |
| 152 | except AttributeError: |
| 153 | return "" |
| 154 | |
| 155 | def _validate_allow_from(self) -> None: |
| 156 | for name, ch in self.channels.items(): |
| 157 | cfg = ch.config |
| 158 | if isinstance(cfg, dict): |
| 159 | if "allow_from" in cfg: |
| 160 | allow = cfg.get("allow_from") |
| 161 | else: |
| 162 | allow = cfg.get("allowFrom") |
| 163 | else: |
| 164 | allow = getattr(cfg, "allow_from", None) |
| 165 | if allow == []: |
| 166 | raise SystemExit( |
| 167 | f'Error: "{name}" has empty allowFrom (denies all). ' |
| 168 | f'Set ["*"] to allow everyone, or add specific user IDs.' |
| 169 | ) |
| 170 | |
| 171 | async def _start_channel(self, name: str, channel: BaseChannel) -> None: |
| 172 | """Start a channel and log any exceptions.""" |
| 173 | try: |
| 174 | await channel.start() |
| 175 | except Exception as e: |
| 176 | logger.error("Failed to start channel {}: {}", name, e) |
| 177 | |
| 178 | async def start_all(self) -> None: |
| 179 | """Start all channels and the outbound dispatcher.""" |
| 180 | if not self.channels: |
| 181 | logger.warning("No channels enabled") |
| 182 | return |
| 183 | |
| 184 | # Start outbound dispatcher |
| 185 | self._dispatch_task = asyncio.create_task(self._dispatch_outbound()) |
| 186 | |
| 187 | # Start channels |
| 188 | tasks = [] |
| 189 | for name, channel in self.channels.items(): |
| 190 | logger.info("Starting {} channel...", name) |
| 191 | tasks.append(asyncio.create_task(self._start_channel(name, channel))) |
| 192 | |
| 193 | self._notify_restart_done_if_needed() |
| 194 | |
| 195 | # Wait for all to complete (they should run forever) |
| 196 | await asyncio.gather(*tasks, return_exceptions=True) |
| 197 | |
| 198 | def _notify_restart_done_if_needed(self) -> None: |
| 199 | """Send restart completion message when runtime env markers are present.""" |
| 200 | notice = consume_restart_notice_from_env() |
| 201 | if not notice: |
| 202 | return |
| 203 | target = self.channels.get(notice.channel) |
| 204 | if not target: |
| 205 | return |
| 206 | asyncio.create_task(self._send_with_retry( |
| 207 | target, |
| 208 | OutboundMessage( |
| 209 | channel=notice.channel, |
| 210 | chat_id=notice.chat_id, |
| 211 | content=format_restart_completed_message(notice.started_at_raw), |
| 212 | ), |
| 213 | )) |
| 214 | |
| 215 | async def stop_all(self) -> None: |
| 216 | """Stop all channels and the dispatcher.""" |
| 217 | logger.info("Stopping all channels...") |
| 218 | |
| 219 | # Stop dispatcher |
| 220 | if self._dispatch_task: |
| 221 | self._dispatch_task.cancel() |
| 222 | try: |
| 223 | await self._dispatch_task |
| 224 | except asyncio.CancelledError: |
| 225 | pass |
| 226 | |
| 227 | # Stop all channels |
| 228 | for name, channel in self.channels.items(): |
| 229 | try: |
| 230 | await channel.stop() |
| 231 | logger.info("Stopped {} channel", name) |
| 232 | except Exception as e: |
| 233 | logger.error("Error stopping {}: {}", name, e) |
| 234 | |
| 235 | async def _dispatch_outbound(self) -> None: |
| 236 | """Dispatch outbound messages to the appropriate channel.""" |
| 237 | logger.info("Outbound dispatcher started") |
| 238 | |
| 239 | # Buffer for messages that couldn't be processed during delta coalescing |
| 240 | # (since asyncio.Queue doesn't support push_front) |
| 241 | pending: list[OutboundMessage] = [] |
| 242 | |
| 243 | while True: |
| 244 | try: |
| 245 | # First check pending buffer before waiting on queue |
| 246 | if pending: |
| 247 | msg = pending.pop(0) |
| 248 | else: |
| 249 | msg = await asyncio.wait_for( |
| 250 | self.bus.consume_outbound(), |
| 251 | timeout=1.0 |
| 252 | ) |
| 253 | |
| 254 | if msg.metadata.get("_progress"): |
| 255 | if msg.metadata.get("_tool_hint") and not self.config.channels.send_tool_hints: |
| 256 | continue |
| 257 | if not msg.metadata.get("_tool_hint") and not self.config.channels.send_progress: |
| 258 | continue |
| 259 | |
| 260 | if msg.metadata.get("_retry_wait"): |
| 261 | continue |
| 262 | |
| 263 | # Coalesce consecutive _stream_delta messages for the same (channel, chat_id) |
| 264 | # to reduce API calls and improve streaming latency |
| 265 | if msg.metadata.get("_stream_delta") and not msg.metadata.get("_stream_end"): |
| 266 | msg, extra_pending = self._coalesce_stream_deltas(msg) |
| 267 | pending.extend(extra_pending) |
| 268 | |
| 269 | channel = self.channels.get(msg.channel) |
| 270 | if channel: |
| 271 | await self._send_with_retry(channel, msg) |
| 272 | else: |
| 273 | logger.warning("Unknown channel: {}", msg.channel) |
| 274 | |
| 275 | except asyncio.TimeoutError: |
| 276 | continue |
| 277 | except asyncio.CancelledError: |
| 278 | break |
| 279 | |
| 280 | @staticmethod |
| 281 | async def _send_once(channel: BaseChannel, msg: OutboundMessage) -> None: |
| 282 | """Send one outbound message without retry policy.""" |
| 283 | if msg.metadata.get("_stream_delta") or msg.metadata.get("_stream_end"): |
| 284 | await channel.send_delta(msg.chat_id, msg.content, msg.metadata) |
| 285 | elif not msg.metadata.get("_streamed"): |
| 286 | await channel.send(msg) |
| 287 | |
| 288 | def _coalesce_stream_deltas( |
| 289 | self, first_msg: OutboundMessage |
| 290 | ) -> tuple[OutboundMessage, list[OutboundMessage]]: |
| 291 | """Merge consecutive _stream_delta messages for the same (channel, chat_id). |
| 292 | |
| 293 | This reduces the number of API calls when the queue has accumulated multiple |
| 294 | deltas, which happens when LLM generates faster than the channel can process. |
| 295 | |
| 296 | Returns: |
| 297 | tuple of (merged_message, list_of_non_matching_messages) |
| 298 | """ |
| 299 | target_key = (first_msg.channel, first_msg.chat_id) |
| 300 | combined_content = first_msg.content |
| 301 | final_metadata = dict(first_msg.metadata or {}) |
| 302 | non_matching: list[OutboundMessage] = [] |
| 303 | |
| 304 | # Only merge consecutive deltas. As soon as we hit any other message, |
| 305 | # stop and hand that boundary back to the dispatcher via `pending`. |
| 306 | while True: |
| 307 | try: |
| 308 | next_msg = self.bus.outbound.get_nowait() |
| 309 | except asyncio.QueueEmpty: |
| 310 | break |
| 311 | |
| 312 | # Check if this message belongs to the same stream |
| 313 | same_target = (next_msg.channel, next_msg.chat_id) == target_key |
| 314 | is_delta = next_msg.metadata and next_msg.metadata.get("_stream_delta") |
| 315 | is_end = next_msg.metadata and next_msg.metadata.get("_stream_end") |
| 316 | |
| 317 | if same_target and is_delta and not final_metadata.get("_stream_end"): |
| 318 | # Accumulate content |
| 319 | combined_content += next_msg.content |
| 320 | # If we see _stream_end, remember it and stop coalescing this stream |
| 321 | if is_end: |
| 322 | final_metadata["_stream_end"] = True |
| 323 | # Stream ended - stop coalescing this stream |
| 324 | break |
| 325 | else: |
| 326 | # First non-matching message defines the coalescing boundary. |
| 327 | non_matching.append(next_msg) |
| 328 | break |
| 329 | |
| 330 | merged = OutboundMessage( |
| 331 | channel=first_msg.channel, |
| 332 | chat_id=first_msg.chat_id, |
| 333 | content=combined_content, |
| 334 | metadata=final_metadata, |
| 335 | ) |
| 336 | return merged, non_matching |
| 337 | |
| 338 | async def _send_with_retry(self, channel: BaseChannel, msg: OutboundMessage) -> None: |
| 339 | """Send a message with retry on failure using exponential backoff. |
| 340 | |
| 341 | Note: CancelledError is re-raised to allow graceful shutdown. |
| 342 | """ |
| 343 | max_attempts = max(self.config.channels.send_max_retries, 1) |
| 344 | |
| 345 | for attempt in range(max_attempts): |
| 346 | try: |
| 347 | await self._send_once(channel, msg) |
| 348 | return # Send succeeded |
| 349 | except asyncio.CancelledError: |
| 350 | raise # Propagate cancellation for graceful shutdown |
| 351 | except Exception as e: |
| 352 | if attempt == max_attempts - 1: |
| 353 | logger.error( |
| 354 | "Failed to send to {} after {} attempts: {} - {}", |
| 355 | msg.channel, max_attempts, type(e).__name__, e |
| 356 | ) |
| 357 | return |
| 358 | delay = _SEND_RETRY_DELAYS[min(attempt, len(_SEND_RETRY_DELAYS) - 1)] |
| 359 | logger.warning( |
| 360 | "Send to {} failed (attempt {}/{}): {}, retrying in {}s", |
| 361 | msg.channel, attempt + 1, max_attempts, type(e).__name__, delay |
| 362 | ) |
| 363 | try: |
| 364 | await asyncio.sleep(delay) |
| 365 | except asyncio.CancelledError: |
| 366 | raise # Propagate cancellation during sleep |
| 367 | |
| 368 | def get_channel(self, name: str) -> BaseChannel | None: |
| 369 | """Get a channel by name.""" |
| 370 | return self.channels.get(name) |
| 371 | |
| 372 | def get_status(self) -> dict[str, Any]: |
| 373 | """Get status of all channels.""" |
| 374 | return { |
| 375 | name: { |
| 376 | "enabled": True, |
| 377 | "running": channel.is_running |
| 378 | } |
| 379 | for name, channel in self.channels.items() |
| 380 | } |
| 381 | |
| 382 | @property |
| 383 | def enabled_channels(self) -> list[str]: |
| 384 | """Get list of enabled channel names.""" |
| 385 | return list(self.channels.keys()) |
| 386 |