agent.py

# agent.py
from __future__ import annotations

import asyncio
import logging
import os
import re
import shutil
import signal
import sys
from pathlib import Path

import yaml
from dotenv import load_dotenv
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler, FileModifiedEvent, FileCreatedEvent, FileMovedEvent

import sqlite3

from core.event_bus import EventBus, Event
from core.health import HealthServer
from core.llm import LLMClient
from core.conversation import ConversationManager
from core.notifier import Notifier
from core.fcm import FcmNotifier
from core.rate_limiter import RateLimiter
from core.scheduler import Scheduler
from core.topic_session import TopicSessionStore
from core.trace import TraceLogFilter, set_trace
from channels.telegram import TelegramChannel
from handlers.message_handler import MessageHandler, CHAT_SYSTEM_PROMPT
from handlers.intent_router import IntentRouter
from tasks.code_task import CodeTaskExecutor
from tasks.execution_tasks import ExecutionTaskStore
from tasks.schema_migration import migrate_todos_schema, migrate_scheduled_messages_schema
from tasks.sync_actor import SyncTodosSchedulesActor, check_offline_catchup
from tasks.sync_client import SyncClient
from tasks.sync_state import SyncState
from tasks.todo import TodoManager
from tasks.scheduled_message import ScheduledMessageManager
from tasks.push_tokens import PushTokenManager
from tasks.cron_dispatcher import CronDispatcher, TodoReminderActor, ScheduledMessageActor, DailyDigestActor, MemoryConsolidationActor
from tasks.result_blog import ResultBlogActor
from channels.web import WebChannel
from core.users import UserManager

BASE_DIR = Path(__file__).parent


class BlogFileHandler(FileSystemEventHandler):
    def __init__(self, event_bus: EventBus) -> None:
        self._bus = event_bus

    def _publish(self, path: str, is_directory: bool) -> None:
        if not is_directory:
            self._bus.publish(Event(type="file_changed", data={"path": path}))

    def on_modified(self, event: FileModifiedEvent) -> None:
        self._publish(event.src_path, event.is_directory)

    def on_created(self, event: FileCreatedEvent) -> None:
        self._publish(event.src_path, event.is_directory)

    def on_moved(self, event: FileMovedEvent) -> None:
        self._publish(event.dest_path, event.is_directory)


def _resolve_env_var(match: re.Match) -> str:
    var_expr = match.group(1)
    if ":-" in var_expr:
        var_name, default = var_expr.split(":-", 1)
    else:
        var_name, default = var_expr, ""
    return os.environ.get(var_name, default)


CREDENTIAL_MAP = {
    "deepseek-api-key": "DEEPSEEK_API_KEY",
    "claude-api-key": "CLAUDE_API_KEY",
    "telegram-bot-token": "TELEGRAM_BOT_TOKEN",
    "web-api-password": "WEB_API_PASSWORD",
    "sync-peer-password": "SYNC_PEER_PASSWORD",
}


def _load_credentials() -> None:
    creds_dir = os.environ.get("CREDENTIALS_DIRECTORY", "")
    if not creds_dir:
        return
    for cred_name, env_var in CREDENTIAL_MAP.items():
        cred_path = Path(creds_dir) / cred_name
        if cred_path.exists():
            os.environ[env_var] = cred_path.read_text().strip()


def load_config() -> dict:
    load_dotenv(BASE_DIR / ".env")
    config_path = BASE_DIR / "config.yaml"
    with open(config_path) as f:
        raw = f.read()
    prev = None
    while prev != raw:
        prev = raw
        raw = re.sub(r"\$\{([^{}]*)\}", _resolve_env_var, raw)
    return yaml.safe_load(raw)


def _blog_repo_root(blog_path: Path) -> Path:
    if blog_path.name == "posts" and blog_path.parent.name == "content" and blog_path.parent.parent.name == "src":
        return blog_path.parent.parent.parent
    return blog_path


def _existing_auto_ro_paths(blog_posts_dir: str) -> list[str]:
    paths = [
        Path.home() / ".agents",
        Path.home() / "agent-ai",
        Path.home() / ".claude" / "skills" / "blog-update",
        # ro so sandboxed pi can git pull/push (ssh key, known_hosts, identity)
        Path.home() / ".ssh",
        Path.home() / ".gitconfig",
    ]
    if blog_posts_dir:
        blog_path = Path(blog_posts_dir).expanduser()
        if blog_path.exists():
            paths.append(_blog_repo_root(blog_path.resolve()))

    out: list[str] = []
    seen: set[str] = set()
    for path in paths:
        if path.exists():
            resolved = str(path.resolve())
            if resolved not in seen:
                seen.add(resolved)
                out.append(resolved)
    return out


def setup_logging(config: dict) -> None:
    from core.chat_log import ChatLogHandler

    log_cfg = config.get("logging", {})
    log_file = BASE_DIR / log_cfg.get("file", "data/logs/agent.log")
    log_file.parent.mkdir(parents=True, exist_ok=True)
    trace_filter = TraceLogFilter()
    handlers = [
        logging.FileHandler(log_file),
        logging.StreamHandler(sys.stdout),
    ]
    for h in handlers:
        h.addFilter(trace_filter)
    logging.basicConfig(
        level=getattr(logging, log_cfg.get("level", "INFO")),
        format="%(asctime)s [%(levelname)s] %(name)s [trace:%(trace)s]: %(message)s",
        handlers=handlers,
    )
    logging.getLogger("core.conversation").addHandler(
        ChatLogHandler(BASE_DIR / "data" / "logs")
    )


async def main() -> None:
    _load_credentials()
    config = load_config()
    setup_logging(config)
    logger = logging.getLogger("agent")

    event_bus = EventBus()
    llm_client = LLMClient(config.get("llm", {}))

    # Conversation context (LangGraph checkpoint)
    conv_cfg = config.get("conversation", {})
    llm_cfg = config.get("llm", {})
    llm_backend = llm_cfg.get("backend", "deepseek")
    model_name = llm_cfg.get("backends", {}).get(llm_backend, {}).get("model", "")
    conversation = ConversationManager(
        db_path=":memory:",
        chat_model=llm_client.get_chat_model(),
        max_context_tokens=conv_cfg.get("max_context_tokens", 4096),
        session_timeout_hours=conv_cfg.get("session_timeout_hours", 24),
        system_prompt=CHAT_SYSTEM_PROMPT,
        memory_db_path=str(BASE_DIR / "data" / "memories.db"),
        model_name=model_name,
    )
    await conversation.start()

    # Database
    local_origin = int(config.get("local_origin", 0))
    sync_node_id = str(config.get("sync", {}).get("node_id", "") or f"origin-{local_origin}")
    db_path = BASE_DIR / "data" / "todos.db"
    db_path.parent.mkdir(parents=True, exist_ok=True)
    db_conn = sqlite3.connect(str(db_path))
    db_conn.execute(
        "CREATE TABLE IF NOT EXISTS todos "
        "(id INTEGER NOT NULL, origin INTEGER NOT NULL, "
        "content TEXT NOT NULL, due_at TEXT, done INTEGER DEFAULT 0, "
        "created_at TEXT NOT NULL, updated_at TEXT NOT NULL, "
        "updated_by TEXT NOT NULL DEFAULT 'legacy', "
        "PRIMARY KEY (origin, id))"
    )
    db_conn.commit()
    migrate_todos_schema(db_conn)
    db_conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_todos_updated_at ON todos(updated_at)"
    )
    db_conn.commit()
    todo_manager = TodoManager(db_conn, default_origin=local_origin, sync_node_id=sync_node_id)

    # Scheduled messages
    db_conn.execute(
        "CREATE TABLE IF NOT EXISTS scheduled_messages "
        "(id INTEGER NOT NULL, origin INTEGER NOT NULL, "
        "content TEXT NOT NULL, action TEXT NOT NULL DEFAULT 'notify', "
        "trigger_at TEXT, cron_expr TEXT, enabled INTEGER DEFAULT 1, "
        "created_at TEXT NOT NULL, updated_at TEXT NOT NULL, "
        "updated_by TEXT NOT NULL DEFAULT 'legacy', "
        "PRIMARY KEY (origin, id))"
    )
    db_conn.commit()
    migrate_scheduled_messages_schema(db_conn)
    db_conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_scheduled_messages_updated_at "
        "ON scheduled_messages(updated_at)"
    )
    db_conn.commit()
    schedule_manager = ScheduledMessageManager(db_conn, default_origin=local_origin, sync_node_id=sync_node_id)
    db_conn.execute(
        "CREATE TABLE IF NOT EXISTS push_tokens "
        "(token TEXT PRIMARY KEY, platform TEXT NOT NULL, updated_at TEXT DEFAULT CURRENT_TIMESTAMP)"
    )
    db_conn.commit()
    push_tokens = PushTokenManager(db_conn)
    credentials_dir = os.environ.get("CREDENTIALS_DIRECTORY", "")
    fcm_credential = Path(credentials_dir) / "fcm-service-account" if credentials_dir else None
    fcm_notifier = FcmNotifier(push_tokens, fcm_credential)

    async def todo_callback(command: str, args: str) -> str:
        parts = args.strip().split(None, 1)
        subcmd = parts[0] if parts else "list"
        rest = parts[1] if len(parts) > 1 else ""

        if subcmd == "list":
            return todo_manager.format_list(todo_manager.list_pending())
        elif subcmd == "add":
            if not rest:
                return "Usage: /todo add <text> [due:YYYY-MM-DD HH:MM]"
            due_at = None
            due_match = re.search(r"\bdue:(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2})\s*$", rest)
            if due_match:
                due_at = due_match.group(1)
                rest = rest[: due_match.start()].strip()
            todo_manager.add(rest, due_at=due_at)
            return f"Added.{f' (due: {due_at})' if due_at else ''}"
        elif subcmd == "done":
            if not rest.isdigit():
                return "Usage: /todo done <id>"
            if todo_manager.done(origin=todo_manager.default_origin, todo_id=int(rest)):
                return f"Todo #{rest} done."
            return f"Todo #{rest} not found."
        else:
            return "Usage: /todo [list|add <text>|done <id>]"

    SCHEDULE_ADD_USAGE = 'Usage: /schedule add <content> at:YYYY-MM-DD HH:MM | cron:"expr" [action:notify|query]'
    SCHEDULE_USAGE = 'Usage: /schedule [list|add <content> at:...|cron:"..." [action:...]|remove <id>]'

    async def schedule_callback(command: str, args: str) -> str:
        parts = args.strip().split(None, 1)
        subcmd = parts[0] if parts else "list"
        rest = parts[1] if len(parts) > 1 else ""

        if subcmd == "list":
            return schedule_manager.format_list(schedule_manager.list_enabled())
        elif subcmd == "add":
            if not rest:
                return SCHEDULE_ADD_USAGE
            trigger_at = None
            cron_expr = None
            action = "notify"
            at_match = re.search(r"\bat[:\s]+(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2})\b", rest)
            cron_match = re.search(r'\bcron:"([^"]+)"', rest)
            action_match = re.search(r"\baction:(notify|query)\b", rest)
            if action_match:
                action = action_match.group(1)
                rest = rest[: action_match.start()] + rest[action_match.end():]
                rest = rest.strip()
            if at_match:
                trigger_at = at_match.group(1)
                rest = rest[: at_match.start()].strip()
            elif cron_match:
                cron_expr = cron_match.group(1)
                rest = rest[: cron_match.start()].strip()
            else:
                return SCHEDULE_ADD_USAGE
            try:
                msg_id = schedule_manager.add(rest, action=action, trigger_at=trigger_at, cron_expr=cron_expr)
                schedule = trigger_at or f"cron:{cron_expr}"
                return f"Scheduled #{msg_id}: {rest} ({schedule}, {action})"
            except ValueError as e:
                return str(e)
        elif subcmd == "remove":
            if not rest.isdigit():
                return "Usage: /schedule remove <id>"
            if schedule_manager.remove(origin=schedule_manager.default_origin, msg_id=int(rest)):
                return f"Scheduled #{rest} removed."
            return f"Scheduled #{rest} not found."
        else:
            return SCHEDULE_USAGE

    # Telegram
    telegram_bot = None
    tg_channel = None
    tg_cfg = config.get("telegram", {})
    tg_token = os.environ.get("TELEGRAM_BOT_TOKEN", "")
    tg_chat_id = os.environ.get("TELEGRAM_CHAT_ID", "")

    rl_cfg = config.get("rate_limit", {})
    msg_limiter = RateLimiter(max_requests=rl_cfg.get("messages_per_minute", 20), window_seconds=60)
    task_limiter = RateLimiter(max_requests=rl_cfg.get("tasks_per_minute", 5), window_seconds=60)

    async def combined_callback(command: str, args: str) -> str:
        if command == "todo":
            return await todo_callback(command, args)
        elif command == "schedule":
            return await schedule_callback(command, args)
        return "Unknown command."

    if tg_cfg.get("enabled") and tg_token and tg_chat_id:
        tg_channel = TelegramChannel(
            token=tg_token,
            chat_id=tg_chat_id,
            event_bus=event_bus,
            command_callback=combined_callback,
            message_limiter=msg_limiter,
            task_limiter=task_limiter,
        )
        max_retries = tg_cfg.get("start_retries", 10)
        retry_delay = tg_cfg.get("start_retry_delay", 30)
        for attempt in range(1, max_retries + 1):
            try:
                await tg_channel.start()
                telegram_bot = tg_channel.app.bot
                logger.info("Telegram connected")
                break
            except Exception as e:
                logger.warning("Telegram start failed (attempt %d/%d): %s", attempt, max_retries, e)
                if attempt == max_retries:
                    logger.error("Telegram start failed after %d retries, giving up", max_retries)
                    raise
                await asyncio.sleep(retry_delay)

    notifier = Notifier(
        config=config.get("notifications", {}),
        telegram_bot=telegram_bot,
        chat_id=tg_chat_id,
    )

    # Message handling (chat + task execution)
    auth_cfg = config.get("auth", {})
    task_backend = config.get("task_backend", "pi")
    pi_workspace_cfg = config.get("pi_workspace", "./pi-workspace")
    pi_workspace_path = (Path(__file__).resolve().parent / pi_workspace_cfg).resolve() if pi_workspace_cfg else None
    if pi_workspace_path:
        pi_workspace_path.mkdir(parents=True, exist_ok=True)
    pi_path = config.get("pi_path", "/home/user/.local/bin/pi")
    blog_posts_dir = config.get("scheduler", {}).get("tasks", {}).get("blog_sync", {}).get("watch", "")

    # pi subprocess sandbox (bubblewrap + env allowlist). See
    # docs/superpowers/specs/2026-07-28-pi-sandbox-design.md
    pi_sandbox_cfg = config.get("pi_sandbox", {}) or {}
    sandbox_ro_paths = list(pi_sandbox_cfg.get("ro_paths") or [])
    for path in _existing_auto_ro_paths(blog_posts_dir):
        if path not in sandbox_ro_paths:
            sandbox_ro_paths.append(path)
    # No blog rw-binds: ResultBlogActor writes posts in this process, never from
    # inside sandboxed pi, so pi needs no write access to the blog repo. The repo
    # stays ro-bound (pi may read it for context). Note _build_bwrap_argv ignores
    # rw_paths entirely whenever a per-user workspace is set.
    sandbox_rw_paths = list(pi_sandbox_cfg.get("rw_paths") or [])
    sandbox_config = {
        "enabled": pi_sandbox_cfg.get("enabled", True),
        "project_dir": pi_sandbox_cfg.get("project_dir") or None,
        "pass_env": pi_sandbox_cfg.get("pass_env") or None,  # None => use _DEFAULT_PASS_ENV
        "ro_paths": sandbox_ro_paths or None,
        "rw_paths": sandbox_rw_paths,
        "fs_enabled": False,  # set below after probe
    }
    logger.info(
        "Pi sandbox config: enabled=%s ro_paths=%s rw_paths=%s",
        sandbox_config["enabled"],
        sandbox_config["ro_paths"] or [],
        sandbox_config["rw_paths"],
    )
    if sandbox_config["enabled"]:
        from tasks.pi_rpc import _probe_bwrap
        fs_ok = _probe_bwrap()
        sandbox_config["fs_enabled"] = fs_ok
        if not fs_ok:
            logger.error("pi_sandbox.enabled but bwrap cannot create a sandbox on this host (unprivileged userns blocked?); pi will run with env allowlist only (no FS sandbox)")
            await notifier.send("⚠️ pi 沙箱: bwrap 在本机无法创建沙箱（userns 受限），仅启用 env 白名单（无 FS 隔离）", level="warning")
            if auth_cfg.get("enabled"):
                raise RuntimeError("multi-tenant mode requires a working bwrap filesystem sandbox")

    code_task_executor = CodeTaskExecutor(
        llm_client, notifier,
        backend=task_backend, pi_workspace=pi_workspace_path, pi_path=pi_path,
        sandbox_config=sandbox_config,
    )

    blog_cfg = config.get("blog", {}) or {}
    blog_actor = (
        ResultBlogActor(
            llm_client,
            blog_posts_dir,
            trace_max_steps=blog_cfg.get("trace_max_steps"),
        )
        if blog_posts_dir else None
    )

    # One store shared by both callers, so /topic/start and a /message
    # "/query_end" for the same caller see the same open session.
    topic_sessions = TopicSessionStore()

    message_handler = MessageHandler(
        event_bus=event_bus,
        conversation=conversation,
        code_task_executor=code_task_executor,
        bot=telegram_bot,
        command_callback=combined_callback,
        intent_router=IntentRouter(llm_client, intent_config=config.get("llm", {})),
        blog_actor=blog_actor,
        topic_sessions=topic_sessions,
    )
    message_handler.start()

    # Multi-tenant identity (cloud deployments)
    user_manager: UserManager | None = None
    if auth_cfg.get("enabled"):
        users_conn = sqlite3.connect(str(BASE_DIR / "data" / "users.db"))
        user_manager = UserManager(users_conn, session_ttl=auth_cfg.get("session_ttl", 86400))
        user_manager.init_schema()
        if auth_cfg.get("admin_email") and auth_cfg.get("admin_password"):
            if user_manager.bootstrap_admin(auth_cfg["admin_email"], auth_cfg["admin_password"]):
                logger.info("Bootstrapped admin user %s", auth_cfg["admin_email"])

    # Web API channel
    web_channel = None
    web_cfg = config.get("web", {})
    quota_cfg = config.get("quota", {})
    if web_cfg.get("enabled") and (web_cfg.get("password") or user_manager is not None):
        execution_db = sqlite3.connect(
            str(BASE_DIR / "data" / "execution_tasks.db"), check_same_thread=False,
        )
        execution_tasks = ExecutionTaskStore(execution_db)
        web_host = "0.0.0.0" if web_cfg.get("bind_public", False) else "127.0.0.1"
        web_channel = WebChannel(
            host=web_host,
            port=web_cfg.get("port", 8765),
            password=web_cfg.get("password", ""),
            allowed_origin=web_cfg.get("allowed_origin", ""),
            conversation=conversation,
            code_task_executor=code_task_executor,
            todo_manager=todo_manager,
            schedule_manager=schedule_manager,
            push_token_manager=push_tokens,
            blog_actor=blog_actor,
            message_handler=message_handler,
            topic_sessions=topic_sessions,
            silent_errors=web_cfg.get("silent_errors", False),
            user_manager=user_manager,
            quota_config=quota_cfg,
            execution_task_store=execution_tasks,
        )
        await web_channel.start()
    dispatcher = CronDispatcher()
    dispatcher.register("todo_reminder", TodoReminderActor(todo_manager, notifier, fcm_notifier, sync_client=None, push_tokens=push_tokens))
    dispatcher.register("scheduled_message", ScheduledMessageActor(schedule_manager, notifier, code_task_executor, telegram_bot, tg_chat_id, fcm_notifier, sync_client=None, push_tokens=push_tokens))
    dispatcher.register("memory_consolidation", MemoryConsolidationActor(conversation))
    dispatcher.register("daily_digest", DailyDigestActor(code_task_executor, telegram_bot, tg_chat_id))

    # Sync (todos + schedules bidirectional)
    sync_cfg = config.get("sync", {})
    sync_state = SyncState(BASE_DIR / "data" / "sync_state.json")
    sync_client = None
    if sync_cfg.get("enabled") and sync_cfg.get("peer_url"):
        sync_client = SyncClient(
            peer_url=sync_cfg["peer_url"],
            peer_password=sync_cfg.get("peer_password", ""),
        )
        sync_actor = SyncTodosSchedulesActor(sync_client, todo_manager, schedule_manager, sync_state)
        dispatcher.register("sync_todos_schedules", sync_actor)
        dispatcher.register("sync_todos_schedules_evening", sync_actor)
        dispatcher.register("todo_reminder", TodoReminderActor(todo_manager, notifier, fcm_notifier, sync_client, push_tokens))
        dispatcher.register("scheduled_message", ScheduledMessageActor(schedule_manager, notifier, code_task_executor, telegram_bot, tg_chat_id, fcm_notifier, sync_client, push_tokens))

    async def _on_cron_tick(event: Event) -> None:
        set_trace()
        await dispatcher.dispatch(event.data)

    loop = asyncio.get_running_loop()

    def _on_cron_tick_sync(event: Event) -> None:
        asyncio.ensure_future(_on_cron_tick(event), loop=loop)

    event_bus.subscribe("cron_tick", _on_cron_tick_sync)

    sched_cfg = config.get("scheduler", {})
    scheduler = None
    if sched_cfg.get("enabled", True):
        scheduler = Scheduler(sched_cfg, event_bus)
        scheduler.start()

    # Health check server
    health_cfg = config.get("health", {})
    health_server = HealthServer(
        host=health_cfg.get("host", "127.0.0.1"),
        port=health_cfg.get("port", 9876),
    )
    health_server.set_check("telegram", {"polling": tg_channel.is_polling_alive if tg_channel else False})
    health_server.set_check("scheduler", {"running": scheduler is not None})
    health_server.set_check("web", {"listening": web_channel is not None})
    health_server.set_check("db", {"ok": True})
    if health_cfg.get("enabled", True):
        await health_server.start()

    # Offline catchup: if sync enabled and last pull > 25h ago, run sync once after 60s
    if sync_client is not None:
        async def _run_catchup() -> None:
            set_trace()
            await asyncio.sleep(60)
            await check_offline_catchup(sync_client, todo_manager, schedule_manager, sync_state)
        asyncio.ensure_future(_run_catchup(), loop=loop)

    # Blog sync task
    observer = None
    if blog_posts_dir:
        from tasks.blog_sync import BlogSyncTask
        _blog_cfg = config.get("scheduler", {}).get("tasks", {}).get("blog_sync", {})
        debounce = _blog_cfg.get("debounce", 1200)
        retry_interval = _blog_cfg.get("retry_interval", 600)
        retry_max = _blog_cfg.get("retry_max_attempts", 3)
        blog_task = BlogSyncTask(
            event_bus, notifier, blog_posts_dir, debounce_seconds=debounce,
            retry_interval_seconds=retry_interval, retry_max_attempts=retry_max,
        )
        blog_task.start()

        handler = BlogFileHandler(event_bus)
        observer = Observer()
        observer.schedule(handler, blog_posts_dir, recursive=True)
        observer.start()
        logger.info("File watcher started on %s", blog_posts_dir)

    await notifier.send("MiniAgent started", level="info")

    # Wait for shutdown
    stop_event = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, stop_event.set)

    await stop_event.wait()

    logger.info("Shutting down...")
    message_handler.stop()
    await conversation.stop()
    if scheduler is not None:
        scheduler.shutdown()
    await health_server.stop()
    if observer:
        observer.stop()
        observer.join()
    await notifier.send("MiniAgent stopped", level="warning")
    if web_channel:
        await web_channel.stop()
    if tg_channel:
        await tg_channel.stop()


if __name__ == "__main__":
    asyncio.run(main())