main
py 639 lines 20.7 KB
Raw
1 import json
2 import os
3 import threading
4 import time
5 import uuid
6 from contextlib import asynccontextmanager, suppress
7
8 from aiogram import Bot
9 from aiogram.client.default import DefaultBotProperties
10 from aiogram.enums import ParseMode
11 from aiogram.types import Message as TgMessage, CallbackQuery
12
13 from agent import AgentContext, UserMessage
14 from helpers import plugins, files, projects
15 from helpers import message_queue as mq
16 from helpers import integration_commands
17 from helpers.notification import NotificationManager, NotificationType, NotificationPriority
18 from helpers.persist_chat import save_tmp_chat
19 from helpers.print_style import PrintStyle
20 from helpers.errors import format_error
21 from initialize import initialize_agent
22
23 from plugins._telegram_integration.helpers import telegram_client as tc
24 from plugins._telegram_integration.helpers import command_ui
25 from plugins._telegram_integration.helpers.bot_manager import get_bot
26 from plugins._telegram_integration.helpers.constants import (
27 PLUGIN_NAME,
28 DOWNLOAD_FOLDER,
29 STATE_FILE,
30 CTX_TG_BOT,
31 CTX_TG_BOT_CFG,
32 CTX_TG_CHAT_ID,
33 CTX_TG_CHAT_TYPE,
34 CTX_TG_USER_ID,
35 CTX_TG_USERNAME,
36 CTX_TG_TYPING_STOP,
37 CTX_TG_REPLY_TO,
38 CTX_TG_ATTACHMENTS,
39 CTX_TG_KEYBOARD,
40 )
41
42 # Chat mapping: (bot_name, tg_user_id) → AgentContext ID
43
44 _chat_map_lock = threading.Lock()
45
46
47 def _load_state() -> dict:
48 path = files.get_abs_path(STATE_FILE)
49 if os.path.isfile(path):
50 try:
51 return json.loads(files.read_file(path))
52 except Exception:
53 return {}
54 return {}
55
56
57 def _save_state(state: dict):
58 path = files.get_abs_path(STATE_FILE)
59 files.make_dirs(path)
60 files.write_file(path, json.dumps(state))
61
62
63 def _map_key(bot_name: str, user_id: int, chat_id: int) -> str:
64 return f"{bot_name}:{user_id}:{chat_id}"
65
66
67 def cleanup_old_attachments():
68 """Remove downloaded attachment files older than per-bot max age. 0 = keep forever."""
69 config = plugins.get_plugin_config(PLUGIN_NAME) or {}
70 bots_cfg = config.get("bots") or []
71 total_removed = 0
72 upload_dir = files.get_abs_path(DOWNLOAD_FOLDER)
73 if not os.path.isdir(upload_dir):
74 return
75 for bot_cfg in bots_cfg:
76 bot_name = bot_cfg.get("name", "")
77 if not bot_name:
78 continue
79 max_age_hours = bot_cfg.get("attachment_max_age_hours", 0)
80 if not max_age_hours or max_age_hours <= 0:
81 continue
82 prefix = f"tg_{bot_name}_"
83 cutoff = time.time() - max_age_hours * 3600
84 for name in os.listdir(upload_dir):
85 if not name.startswith(prefix):
86 continue
87 path = os.path.join(upload_dir, name)
88 try:
89 if os.path.isfile(path) and os.path.getmtime(path) < cutoff:
90 os.remove(path)
91 total_removed += 1
92 except OSError:
93 pass
94 if total_removed:
95 PrintStyle.info(f"Telegram: cleaned up {total_removed} old attachment(s)")
96
97 # Access control
98
99 def _is_allowed(bot_cfg: dict, user_id: int, username: str | None) -> bool:
100 allowed = bot_cfg.get("allowed_users") or []
101 if not allowed:
102 return True # empty = allow all
103 for entry in allowed:
104 entry_str = str(entry).strip()
105 if entry_str.startswith("@"):
106 if username and f"@{username}" == entry_str:
107 return True
108 else:
109 try:
110 if int(entry_str) == user_id:
111 return True
112 except ValueError:
113 if username and entry_str.lower() == username.lower():
114 return True
115 return False
116
117
118 def _get_project(bot_cfg: dict, user_id: int) -> str:
119 user_projects = bot_cfg.get("user_projects") or {}
120 project = user_projects.get(str(user_id), "")
121 if not project:
122 project = bot_cfg.get("default_project", "")
123 return project
124
125 # Message handlers (registered with aiogram by bot_manager)
126
127 async def handle_start(message: TgMessage, bot_name: str, bot_cfg: dict):
128 """Handle /start command."""
129 user = message.from_user
130 if not user:
131 return
132
133 if not _is_allowed(bot_cfg, user.id, user.username):
134 await message.reply("You are not authorized to use this bot.")
135 return
136
137 instance = get_bot(bot_name)
138 if not instance:
139 return
140
141 await _send_with_temp_bot(
142 instance.bot.token, message.chat.id,
143 f"\U0001f44b Hello {user.first_name}! I'm connected to Agent Zero.\n\n"
144 "Send me a message and I'll process it.\n"
145 "Use /clear to reset the conversation.\n"
146 "Use /project, /model, /agent, or /send to control the current chat.",
147 parse_mode=None,
148 reply_to_message_id=message.message_id,
149 )
150
151 # Ensure a chat context exists
152 await _get_or_create_context(bot_name, bot_cfg, message)
153
154
155 async def handle_message(message: TgMessage, bot_name: str, bot_cfg: dict):
156 """Handle incoming user message."""
157 user = message.from_user
158 if not user:
159 return
160
161 if not _is_allowed(bot_cfg, user.id, user.username):
162 return
163
164 instance = get_bot(bot_name)
165 if not instance:
166 return
167
168 text = _extract_message_content(message)
169 context = await _get_or_create_context(bot_name, bot_cfg, message)
170 if not context:
171 await _send_with_temp_bot(
172 instance.bot.token, message.chat.id,
173 "Failed to create chat session.",
174 parse_mode=None,
175 )
176 return
177 context.data[CTX_TG_CHAT_TYPE] = str(message.chat.type or "")
178 context.data[CTX_TG_REPLY_TO] = message.message_id
179
180 if await command_ui.handle_command(
181 context,
182 instance.bot.token,
183 message.chat.id,
184 message.message_id,
185 text,
186 ):
187 return
188
189 command_reply = integration_commands.try_handle_command(context, text, integration="telegram")
190 if command_reply is not None:
191 await _send_with_temp_bot(
192 instance.bot.token,
193 message.chat.id,
194 command_reply,
195 parse_mode=None,
196 reply_to_message_id=message.message_id,
197 )
198 return
199 if integration_commands.extract_command_line(text).startswith("/"):
200 command = integration_commands.extract_command_line(text).split(" ", 1)[0]
201 await _send_with_temp_bot(
202 instance.bot.token,
203 message.chat.id,
204 integration_commands.unknown_command_text(command, integration="telegram"),
205 parse_mode=None,
206 reply_to_message_id=message.message_id,
207 )
208 return
209
210 # Use temp bot for downloads (cross-event-loop safe)
211 async with _temp_bot(instance.bot.token) as dl_bot:
212 attachments = await _download_attachments(dl_bot, message, bot_name=bot_name)
213
214 # Build user message with prompt
215 agent = context.agent0
216 user_msg = agent.read_prompt(
217 "fw.telegram.user_message.md",
218 sender=_format_user(user),
219 body=text,
220 )
221
222 if context.is_running():
223 item = mq.add(context, user_msg, attachments)
224 save_tmp_chat(context)
225 await _send_with_temp_bot(
226 instance.bot.token,
227 message.chat.id,
228 f"Queued message #{item.get('seq', len(mq.get_queue(context)))}. Use /send to flush queued work, or /steer <message> to interrupt the active run.",
229 parse_mode=None,
230 reply_to_message_id=message.message_id,
231 )
232 return
233
234 # Start persistent typing indicator (thread-based, works across event loops)
235 typing_stop = _start_typing(instance.bot.token, message.chat.id)
236
237 # Store stop event so send_telegram_reply can cancel typing
238 context.data[CTX_TG_TYPING_STOP] = typing_stop
239
240 msg_id = str(uuid.uuid4())
241 mq.log_user_message(context, user_msg, attachments, message_id=msg_id, source=" (telegram)")
242 context.communicate(UserMessage(
243 message=user_msg,
244 attachments=attachments,
245 id=msg_id,
246 ))
247
248 save_tmp_chat(context)
249
250 # Send notification
251 if bot_cfg.get("notify_messages", False):
252 username_str = f"@{user.username}" if user.username else str(user.id)
253 preview = (text[:80] + "...") if len(text) > 80 else text
254 NotificationManager.send_notification(
255 type=NotificationType.INFO,
256 priority=NotificationPriority.HIGH,
257 title="Telegram: new message",
258 message=f"From {username_str}: {preview}",
259 display_time=10,
260 group="telegram",
261 )
262
263
264 async def handle_callback_query(query: CallbackQuery, bot_name: str, bot_cfg: dict):
265 """Handle inline keyboard button press."""
266 user = query.from_user
267 if not user or not query.message:
268 return
269
270 if not _is_allowed(bot_cfg, user.id, user.username):
271 await query.answer("Not authorized.")
272 return
273
274 await query.answer()
275
276 # Treat callback data as a user message
277 text = query.data or ""
278 if not text:
279 return
280
281 context = await _get_or_create_context_from_user(
282 bot_name, bot_cfg, user.id, user.username, query.message.chat.id, str(query.message.chat.type or ""),
283 )
284 if not context:
285 return
286 context.data[CTX_TG_REPLY_TO] = query.message.message_id
287
288 instance = get_bot(bot_name)
289 if instance:
290 try:
291 if await command_ui.handle_callback(
292 context,
293 instance.bot.token,
294 query.message.chat.id,
295 query.message.message_id,
296 text,
297 ):
298 return
299 except Exception as e:
300 PrintStyle.error(f"Telegram callback failed: {format_error(e)}")
301 if text.startswith("tg:"):
302 return
303 if text.startswith("tg:"):
304 return
305
306 command_reply = integration_commands.try_handle_command(context, text, integration="telegram")
307 if command_reply is not None:
308 if instance:
309 await _send_with_temp_bot(
310 instance.bot.token,
311 query.message.chat.id,
312 command_reply,
313 parse_mode=None,
314 reply_to_message_id=query.message.message_id,
315 )
316 return
317
318 agent = context.agent0
319 user_msg = agent.read_prompt(
320 "fw.telegram.user_message.md",
321 sender=_format_user(user),
322 body=f"[Button pressed: {text}]",
323 )
324
325 msg_id = str(uuid.uuid4())
326 mq.log_user_message(context, user_msg, [], message_id=msg_id, source=" (telegram)")
327 context.communicate(UserMessage(message=user_msg, id=msg_id))
328 save_tmp_chat(context)
329
330
331 async def handle_new_members(message: TgMessage, bot_name: str, bot_cfg: dict):
332 """Send welcome message when new members join a group."""
333 if not bot_cfg.get("welcome_enabled", False):
334 return
335
336 new_members = message.new_chat_members or []
337 if not new_members:
338 return
339
340 instance = get_bot(bot_name)
341 if not instance:
342 return
343
344 template = bot_cfg.get("welcome_message", "").strip()
345 if not template:
346 template = "Welcome, {name}!"
347
348 for member in new_members:
349 if member.is_bot:
350 continue
351 name = member.full_name or member.first_name or str(member.id)
352 text = template.replace("{name}", name)
353 await _send_with_temp_bot(instance.bot.token, message.chat.id, text, parse_mode=None)
354
355 # Context management
356
357 async def _get_or_create_context(
358 bot_name: str,
359 bot_cfg: dict,
360 message: TgMessage,
361 ) -> AgentContext | None:
362 user = message.from_user
363 if not user:
364 return None
365 return await _get_or_create_context_from_user(
366 bot_name, bot_cfg, user.id, user.username, message.chat.id, str(message.chat.type or ""),
367 )
368
369
370 async def _get_or_create_context_from_user(
371 bot_name: str,
372 bot_cfg: dict,
373 user_id: int,
374 username: str | None,
375 chat_id: int,
376 chat_type: str = "",
377 ) -> AgentContext | None:
378 key = _map_key(bot_name, user_id, chat_id)
379
380 with _chat_map_lock:
381 state = _load_state()
382 chats = state.setdefault("chats", {})
383 ctx_id = chats.get(key)
384
385 # Check if existing context is still alive
386 if ctx_id:
387 ctx = AgentContext.get(ctx_id)
388 if ctx:
389 ctx.data[CTX_TG_CHAT_TYPE] = chat_type or ctx.data.get(CTX_TG_CHAT_TYPE, "")
390 return ctx
391 # Context was garbage collected, remove stale mapping
392 chats.pop(key, None)
393
394 # Create new context
395 try:
396 config = initialize_agent()
397 display_name = f"@{username}" if username else str(user_id)
398 ctx = AgentContext(config, name=f"Telegram: {display_name}")
399
400 ctx.data[CTX_TG_BOT] = bot_name
401 ctx.data[CTX_TG_BOT_CFG] = bot_cfg
402 ctx.data[CTX_TG_CHAT_ID] = chat_id
403 ctx.data[CTX_TG_CHAT_TYPE] = chat_type
404 ctx.data[CTX_TG_USER_ID] = user_id
405 ctx.data[CTX_TG_USERNAME] = username or ""
406
407 project = _get_project(bot_cfg, user_id)
408 if project:
409 projects.activate_project(ctx.id, project)
410
411 # Inherit model override from an existing context in the same project
412 _inherit_model_override(ctx)
413
414 chats[key] = ctx.id
415 _save_state(state)
416
417 PrintStyle.success(
418 f"Telegram ({bot_name}): new chat {ctx.id} for user {display_name}"
419 )
420 return ctx
421
422 except Exception as e:
423 PrintStyle.error(f"Telegram: failed to create context: {format_error(e)}")
424 return None
425
426 # Message content extraction
427
428 def _extract_message_content(message: TgMessage) -> str:
429 parts = []
430
431 if message.text:
432 parts.append(message.text)
433 elif message.caption:
434 parts.append(message.caption)
435
436 if message.location:
437 loc = message.location
438 parts.append(f"[Location: {loc.latitude}, {loc.longitude}]")
439
440 if message.contact:
441 c = message.contact
442 parts.append(f"[Contact: {c.first_name} {c.last_name or ''} phone={c.phone_number}]")
443
444 if message.sticker:
445 parts.append(f"[Sticker: {message.sticker.emoji or ''}]")
446
447 # Simple attachment indicators
448 for attr, label in [("voice", "Voice message"), ("video_note", "Video note")]:
449 if getattr(message, attr, None):
450 parts.append(f"[{label} — see attachment]")
451
452 return "\n".join(parts) if parts else "[No text content]"
453
454
455 async def _download_attachments(bot, message: TgMessage, bot_name: str = "") -> list[str]:
456 """Download photos, documents, audio, voice, video from message."""
457 paths: list[str] = []
458 tg_prefix = f"tg_{bot_name}_" if bot_name else "tg_"
459 # Host-local path for actual file I/O
460 download_dir = files.get_abs_path(DOWNLOAD_FOLDER)
461 os.makedirs(download_dir, exist_ok=True)
462 # Docker-style path for agent references
463 download_dir_ref = files.get_abs_path_dockerized(DOWNLOAD_FOLDER)
464
465 async def _dl(file_id: str, filename: str) -> str | None:
466 safe_name = f"{tg_prefix}{uuid.uuid4().hex[:8]}_{filename}"
467 dest = os.path.join(download_dir, safe_name)
468 result = await tc.download_file(bot, file_id, dest)
469 if result:
470 return os.path.join(download_dir_ref, safe_name)
471 return None
472
473 # Photo: get largest resolution
474 if message.photo:
475 photo = message.photo[-1]
476 path = await _dl(photo.file_id, f"photo_{photo.file_unique_id}.jpg")
477 if path:
478 paths.append(path)
479
480 # Other attachment types: (attr, default_prefix, default_ext)
481 _types = [
482 ("document", "file", None),
483 ("audio", "audio", ".mp3"),
484 ("voice", "voice", ".ogg"),
485 ("video", "video", ".mp4"),
486 ("video_note", "videonote", ".mp4"),
487 ]
488 for attr, prefix, ext in _types:
489 obj = getattr(message, attr, None)
490 if not obj:
491 continue
492 fname = getattr(obj, "file_name", None) or f"{prefix}_{obj.file_unique_id}{ext or ''}"
493 path = await _dl(obj.file_id, fname)
494 if path:
495 paths.append(path)
496
497 return paths
498
499 # Reply sending (called from process_chain_end extension)
500
501 async def send_telegram_reply(
502 context: AgentContext,
503 response_text: str,
504 attachments: list[str] | None = None,
505 keyboard: list[list[dict]] | None = None,
506 ) -> str | None:
507 """Send reply to Telegram user. Returns error string or None on success."""
508 bot_name = context.data.get(CTX_TG_BOT)
509 if not bot_name:
510 return "No Telegram bot configured on context"
511
512 instance = get_bot(bot_name)
513 if not instance:
514 return f"Bot '{bot_name}' not running"
515
516 chat_id = context.data.get(CTX_TG_CHAT_ID)
517 if not chat_id:
518 return "No chat_id on context"
519
520 reply_to = context.data.get(CTX_TG_REPLY_TO)
521
522 try:
523 async with _temp_bot(instance.bot.token, default=DefaultBotProperties(parse_mode=ParseMode.HTML)) as reply_bot:
524 if attachments:
525 for path in attachments:
526 local_path = files.fix_dev_path(path)
527 if tc.is_image_file(local_path):
528 await tc.send_photo(reply_bot, chat_id, local_path, reply_to_message_id=reply_to)
529 elif tc.is_voice_file(local_path):
530 await tc.send_voice(reply_bot, chat_id, local_path, reply_to_message_id=reply_to)
531 elif tc.is_audio_file(local_path):
532 await tc.send_audio(reply_bot, chat_id, local_path, reply_to_message_id=reply_to)
533 elif tc.is_video_file(local_path):
534 await tc.send_video(reply_bot, chat_id, local_path, reply_to_message_id=reply_to)
535 else:
536 await tc.send_file(reply_bot, chat_id, local_path, reply_to_message_id=reply_to)
537
538 if response_text:
539 html_text = tc.md_to_telegram_html(response_text)
540 from plugins._telegram_integration.helpers import draft_stream
541
542 if await draft_stream.finalize_response(context, response_text, keyboard):
543 pass
544 elif keyboard:
545 await tc.send_text_with_keyboard(reply_bot, chat_id, html_text, keyboard, reply_to_message_id=reply_to)
546 else:
547 await tc.send_text(reply_bot, chat_id, html_text, reply_to_message_id=reply_to)
548
549 return None
550
551 except Exception as e:
552 error = format_error(e)
553 PrintStyle.error(f"Telegram reply failed: {error}")
554 return error
555
556 # Helpers
557
558 @asynccontextmanager
559 async def _temp_bot(token: str, **kwargs):
560 """Create a temporary Bot, yield it, and ensure the session is closed."""
561 bot = Bot(token=token, **kwargs)
562 try:
563 yield bot
564 finally:
565 with suppress(Exception):
566 await bot.session.close()
567
568
569 async def _send_with_temp_bot(
570 token: str,
571 chat_id: int,
572 text: str,
573 parse_mode: str | None = None,
574 reply_to_message_id: int | None = None,
575 ):
576 """Send text using a temporary Bot to avoid cross-event-loop session issues."""
577 async with _temp_bot(token) as bot:
578 await tc.send_text(
579 bot,
580 chat_id,
581 text,
582 reply_to_message_id=reply_to_message_id,
583 parse_mode=parse_mode,
584 )
585
586
587 def _start_typing(token: str, chat_id: int) -> threading.Event:
588 """Spawn a daemon thread that sends typing every 4s. Returns a stop Event."""
589 stop = threading.Event()
590
591 def _run():
592 import asyncio
593
594 async def _loop():
595 async with _temp_bot(token) as bot:
596 while not stop.is_set():
597 await tc.send_typing(bot, chat_id)
598 for _ in range(8):
599 if stop.is_set():
600 return
601 await asyncio.sleep(0.5)
602
603 try:
604 asyncio.run(_loop())
605 except Exception:
606 pass
607
608 threading.Thread(target=_run, daemon=True).start()
609 return stop
610
611
612 def _format_user(user) -> str:
613 name = user.first_name or ""
614 if user.last_name:
615 name += f" {user.last_name}"
616 if user.username:
617 name += f" (@{user.username})"
618 return name.strip() or str(user.id)
619
620
621 def _inherit_model_override(ctx: AgentContext):
622 """Copy chat_model_override from the most recent sibling context in the same project."""
623 project = ctx.get_data("project")
624 if not project:
625 return
626 try:
627 from plugins._model_config.helpers.model_config import is_chat_override_allowed
628 if not is_chat_override_allowed(ctx.agent0):
629 return
630 except Exception:
631 return
632 source = max(
633 (c for c in AgentContext.all()
634 if c.id != ctx.id and c.get_data("project") == project and c.get_data("chat_model_override")),
635 key=lambda c: c.last_message,
636 default=None,
637 )
638 if source:
639 ctx.set_data("chat_model_override", source.get_data("chat_model_override"))