ForcePilot/backend/package/yuxi/channels/adapters/feishu/dynamic_agent.py

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from typing import Any

from yuxi.utils.logging_config import logger


@dataclass
class FeishuDynamicAgentConfig:
    enabled: bool = False
    max_agent_count: int = 50
    ttl_s: float = 3600.0
    template_agent_id: int = 0
    auto_create: bool = False
    workspace_template: str = ""
    agent_dir_template: str = ""
    config_file_template: str = ""

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> FeishuDynamicAgentConfig:
        dac = config.get("dynamicAgentCreation", {}) or {}
        return cls(
            enabled=dac.get("enabled", False),
            max_agent_count=int(dac.get("maxCount", 50)),
            ttl_s=float(dac.get("ttl", 3600)),
            template_agent_id=int(dac.get("templateAgentId", 0)),
            auto_create=dac.get("autoCreate", False),
            workspace_template=dac.get("workspaceTemplate", ""),
            agent_dir_template=dac.get("agentDirTemplate", ""),
            config_file_template=dac.get("configFileTemplate", ""),
        )


class FeishuDynamicAgentManager:
    def __init__(self, config: FeishuDynamicAgentConfig | None = None):
        self.config = config or FeishuDynamicAgentConfig()
        self._dm_agents: dict[str, str] = {}
        self._agent_metadata: dict[str, dict[str, Any]] = {}
        self._agent_created_at: dict[str, float] = {}
        self._agent_last_access: dict[str, float] = {}

    @property
    def active_count(self) -> int:
        return len(self._dm_agents)

    def get_or_resolve_agent(self, open_id: str) -> str | None:
        if not self.config.enabled:
            return None
        if open_id in self._dm_agents:
            self._agent_last_access[open_id] = time.monotonic()
            return self._dm_agents[open_id]
        if self.config.auto_create:
            agent_id = self._create_dm_agent(open_id)
            if agent_id:
                return agent_id
        return None

    def _create_dm_agent(self, open_id: str) -> str | None:
        if len(self._dm_agents) >= self.config.max_agent_count:
            logger.warning("[DynamicAgent] Max agent count reached (%d)", self.config.max_agent_count)
            self._evict_expired()
            if len(self._dm_agents) >= self.config.max_agent_count:
                return None

        agent_id = f"feishu-{open_id}"
        now = time.monotonic()
        self._dm_agents[open_id] = agent_id
        self._agent_created_at[open_id] = now
        self._agent_last_access[open_id] = now
        self._agent_metadata[open_id] = {
            "open_id": open_id,
            "agent_id": agent_id,
            "created_at": now,
        }
        logger.info("[DynamicAgent] Created DM agent for %s: %s", open_id, agent_id)

        if self.config.workspace_template:
            self._ensure_agent_workspace(open_id, agent_id)

        if self.config.config_file_template:
            self._write_agent_config(open_id, agent_id)

        return agent_id

    def _ensure_agent_workspace(self, open_id: str, agent_id: str) -> None:
        ws_template = self.config.workspace_template
        dir_template = self.config.agent_dir_template or "{agent_id}"

        agent_dir = dir_template.replace("{agent_id}", agent_id).replace("{open_id}", open_id)
        workspace_path = os.path.join(ws_template, agent_dir)

        try:
            os.makedirs(workspace_path, exist_ok=True)
            logger.info("[DynamicAgent] Created agent workspace: %s", workspace_path)
        except OSError as e:
            logger.warning("[DynamicAgent] Failed to create workspace %s: %s", workspace_path, e)

    def _write_agent_config(self, open_id: str, agent_id: str) -> None:
        config_template = self.config.config_file_template
        if not config_template:
            return

        config_path = config_template.replace("{agent_id}", agent_id).replace("{open_id}", open_id)
        config_dir = os.path.dirname(config_path)
        if config_dir:
            try:
                os.makedirs(config_dir, exist_ok=True)
            except OSError:
                pass

        config_data = {
            "agent_id": agent_id,
            "open_id": open_id,
            "dynamic": True,
            "created_at": self._agent_created_at.get(open_id, 0),
        }

        try:
            with open(config_path, "w", encoding="utf-8") as f:
                json.dump(config_data, f, ensure_ascii=False, indent=2)
            logger.info("[DynamicAgent] Wrote agent config: %s", config_path)
        except OSError as e:
            logger.warning("[DynamicAgent] Failed to write config %s: %s", config_path, e)

    def _evict_expired(self) -> int:
        if not self.config.ttl_s:
            return 0

        now = time.monotonic()
        expired = [oid for oid, created in self._agent_created_at.items() if (now - created) > self.config.ttl_s]
        for oid in expired:
            self.remove_agent(oid)
        if expired:
            logger.info("[DynamicAgent] Evicted %d expired agents", len(expired))
        return len(expired)

    def remove_agent(self, open_id: str) -> None:
        self._dm_agents.pop(open_id, None)
        self._agent_created_at.pop(open_id, None)
        self._agent_last_access.pop(open_id, None)
        self._agent_metadata.pop(open_id, None)

    def list_agents(self) -> list[dict[str, Any]]:
        result: list[dict[str, Any]] = []
        now = time.monotonic()
        for open_id, agent_id in self._dm_agents.items():
            metadata = self._agent_metadata.get(open_id, {})
            result.append(
                {
                    "agent_id": agent_id,
                    "open_id": open_id,
                    "created_at": self._agent_created_at.get(open_id, 0),
                    "age_s": round(now - self._agent_created_at.get(open_id, now), 1),
                    "last_access_s": round(now - self._agent_last_access.get(open_id, now), 1),
                    "metadata": metadata,
                }
            )
        return result

    def agent_count(self) -> int:
        return len(self._dm_agents)

    def clear(self) -> None:
        self._dm_agents.clear()
        self._agent_created_at.clear()
        self._agent_last_access.clear()
        self._agent_metadata.clear()