illusion-code 0.1.0__py3-none-any.whl

illusion/swarm/in_process.py

@@ -0,0 +1,259 @@
+"""
+进程内代理执行后端模块
+=====================
+
+本模块实现进程内的代理执行后端。
+使用 :mod:`contextvars` 在当前 Python 进程中将代理作为 asyncio Task 运行，
+实现每个代理的上下文隔离。
+
+主要组件：
+    - InProcessBackend: 进程内执行后端，实现 TeammateExecutor 协议
+
+使用示例：
+    >>> from illusion.swarm.in_process import InProcessBackend
+    >>> backend = InProcessBackend()
+    >>> result = await backend.spawn(config)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+
+from illusion.swarm.agent_executor import (
+    AgentExecutionContext,
+    AgentSpawnConfig,
+    AgentAbortController,
+    set_agent_context,
+    _register_agent,
+    _unregister_agent,
+)
+from illusion.swarm.types import (
+    BackendType,
+    SpawnResult,
+    TeammateMessage,
+    TeammateSpawnConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# InProcessBackend
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _AgentEntry:
+    """运行中进程内代理的内部注册表条目。"""
+
+    task: asyncio.Task[None]
+    abort_controller: AgentAbortController
+    task_id: str
+    started_at: float = field(default_factory=time.time)
+
+
+class InProcessBackend:
+    """将代理作为当前进程中的 asyncio Task 运行的 TeammateExecutor。"""
+
+    type: BackendType = "in_process"
+
+    def __init__(self) -> None:
+        self._active: dict[str, _AgentEntry] = {}
+
+    def is_available(self) -> bool:
+        """进程内后端始终可用。"""
+        return True
+
+    async def spawn(self, config: TeammateSpawnConfig) -> SpawnResult:
+        """将进程内代理生成为 asyncio Task。"""
+        agent_id = f"{config.name}@{config.team}"
+        task_id = f"in_process_{uuid.uuid4().hex[:12]}"
+
+        # 检查是否已存在活跃的同名代理
+        if agent_id in self._active:
+            entry = self._active[agent_id]
+            if not entry.task.done():
+                logger.warning("[InProcessBackend] spawn(): %s is already running", agent_id)
+                return SpawnResult(
+                    task_id=task_id,
+                    agent_id=agent_id,
+                    backend_type=self.type,
+                    success=False,
+                    error=f"Agent {agent_id!r} is already running",
+                )
+
+        # 创建中止控制器
+        abort_controller = AgentAbortController()
+
+        # 创建代理生成配置
+        spawn_config = AgentSpawnConfig(
+            name=config.name,
+            prompt=config.prompt,
+            cwd=config.cwd,
+            model=config.model,
+            system_prompt=config.system_prompt,
+            permission_mode=None,
+            parent_session_id=config.parent_session_id,
+        )
+
+        # 预先创建并注册执行上下文，以便 send_message 可以立即找到代理
+        from illusion.swarm.agent_executor import AgentExecutionContext, _register_agent
+        ctx = AgentExecutionContext(
+            agent_id=agent_id,
+            agent_name=config.name,
+            prompt=config.prompt,
+            model=config.model,
+            cwd=__import__("pathlib").Path(config.cwd),
+            abort_controller=abort_controller,
+        )
+        _register_agent(ctx)
+
+        # 创建 asyncio Task
+        task = asyncio.create_task(
+            self._run_agent(spawn_config, agent_id, abort_controller, ctx),
+            name=f"agent-{agent_id}",
+        )
+
+        entry = _AgentEntry(
+            task=task,
+            abort_controller=abort_controller,
+            task_id=task_id,
+        )
+        self._active[agent_id] = entry
+
+        # 添加完成回调
+        def _on_done(t: asyncio.Task[None]) -> None:
+            self._active.pop(agent_id, None)
+            if not t.cancelled() and t.exception() is not None:
+                logger.error("[InProcessBackend] Agent %s raised exception: %s", agent_id, t.exception())
+
+        task.add_done_callback(_on_done)
+
+        logger.debug("[InProcessBackend] spawned %s (task_id=%s)", agent_id, task_id)
+        return SpawnResult(
+            task_id=task_id,
+            agent_id=agent_id,
+            backend_type=self.type,
+        )
+
+    async def _run_agent(
+        self,
+        config: AgentSpawnConfig,
+        agent_id: str,
+        abort_controller: AgentAbortController,
+        ctx: AgentExecutionContext | None = None,
+    ) -> None:
+        """运行代理的内部协程。"""
+        if ctx is None:
+            ctx = AgentExecutionContext(
+                agent_id=agent_id,
+                agent_name=config.name,
+                prompt=config.prompt,
+                model=config.model,
+                cwd=__import__("pathlib").Path(config.cwd),
+                abort_controller=abort_controller,
+            )
+            _register_agent(ctx)
+        set_agent_context(ctx)
+
+        try:
+            # 注意：这里需要 query_context 和 parent_registry
+            # 在实际使用中，这些应该通过某种方式传入
+            # 目前这是一个占位实现
+            logger.info("[InProcessBackend] %s: agent started (stub)", agent_id)
+            ctx.status = "running"
+
+            # 等待取消或完成
+            while not abort_controller.is_cancelled:
+                await asyncio.sleep(0.1)
+
+        except asyncio.CancelledError:
+            logger.debug("[InProcessBackend] %s: task cancelled", agent_id)
+            raise
+        except Exception:
+            logger.exception("[InProcessBackend] %s: unhandled exception", agent_id)
+        finally:
+            ctx.status = "stopped"
+            _unregister_agent(agent_id)
+
+    async def send_message(self, agent_id: str, message: TeammateMessage) -> None:
+        """向运行中的代理发送消息。"""
+        # 首先检查 self._active 中是否有该代理
+        entry = self._active.get(agent_id)
+        if entry is not None and not entry.task.done():
+            # 代理正在运行，但需要找到其 AgentExecutionContext
+            # 从全局注册表中查找
+            from illusion.swarm.agent_executor import get_active_agent
+            agent_ctx = get_active_agent(agent_id)
+            if agent_ctx is not None:
+                await agent_ctx.message_queue.put(message)
+                logger.debug("[InProcessBackend] sent message to %s", agent_id)
+                return
+
+        # 回退：尝试从全局注册表查找
+        from illusion.swarm.agent_executor import get_active_agent, get_active_agent_by_name
+        agent_name = agent_id.split("@")[0] if "@" in agent_id else agent_id
+
+        agent_ctx = get_active_agent(agent_id)
+        if agent_ctx is None:
+            agent_ctx = get_active_agent_by_name(agent_name)
+
+        if agent_ctx is not None:
+            await agent_ctx.message_queue.put(message)
+            logger.debug("[InProcessBackend] sent message to %s", agent_id)
+        else:
+            raise ValueError(f"No active agent found for {agent_id!r}")
+
+    async def shutdown(self, agent_id: str, *, force: bool = False, timeout: float = 10.0) -> bool:
+        """终止运行中的进程内代理。"""
+        entry = self._active.get(agent_id)
+        if entry is None:
+            logger.debug("[InProcessBackend] shutdown(): %s not found", agent_id)
+            return False
+
+        if entry.task.done():
+            self._active.pop(agent_id, None)
+            return True
+
+        if force:
+            entry.abort_controller.request_cancel(reason="force shutdown", force=True)
+            entry.task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await asyncio.wait_for(asyncio.shield(entry.task), timeout=timeout)
+        else:
+            entry.abort_controller.request_cancel(reason="graceful shutdown")
+            try:
+                await asyncio.wait_for(asyncio.shield(entry.task), timeout=timeout)
+            except asyncio.TimeoutError:
+                logger.warning("[InProcessBackend] %s did not exit within %.1fs — forcing", agent_id, timeout)
+                entry.abort_controller.request_cancel(reason="timeout — forcing", force=True)
+                entry.task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await entry.task
+
+        self._active.pop(agent_id, None)
+        logger.debug("[InProcessBackend] shut down %s", agent_id)
+        return True
+
+    def list_agents(self) -> list[tuple[str, bool, float]]:
+        """返回 (agent_id, is_running, duration_seconds) 元组列表。"""
+        now = time.time()
+        result = []
+        for agent_id, entry in self._active.items():
+            is_running = not entry.task.done()
+            duration = now - entry.started_at
+            result.append((agent_id, is_running, duration))
+        return result
+
+    async def shutdown_all(self, *, force: bool = False, timeout: float = 10.0) -> None:
+        """终止所有活跃代理。"""
+        agent_ids = list(self._active.keys())
+        await asyncio.gather(
+            *(self.shutdown(aid, force=force, timeout=timeout) for aid in agent_ids),
+            return_exceptions=True,
+        )

illusion/swarm/subprocess_backend.py

@@ -0,0 +1,136 @@
+"""
+子进程代理执行后端模块
+=====================
+
+本模块实现基于子进程的 TeammateExecutor 接口。
+使用现有的 :class:`~illusion.tasks.manager.BackgroundTaskManager`
+来创建和管理子进程，通过 stdin/stdout 进行通信。
+
+主要组件：
+    - SubprocessBackend: 子进程执行后端
+
+使用示例：
+    >>> from illusion.swarm.subprocess_backend import SubprocessBackend
+    >>> backend = SubprocessBackend()
+    >>> result = await backend.spawn(config)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+
+from illusion.swarm.agent_executor import (
+    _build_agent_cli_flags,
+    _build_agent_env_vars,
+    _get_agent_command,
+)
+from illusion.swarm.types import (
+    BackendType,
+    SpawnResult,
+    TeammateMessage,
+    TeammateSpawnConfig,
+)
+from illusion.tasks.manager import get_task_manager
+
+logger = logging.getLogger(__name__)
+
+
+class SubprocessBackend:
+    """TeammateExecutor 实现，每个代理作为独立子进程运行。"""
+
+    type: BackendType = "subprocess"
+
+    def __init__(self) -> None:
+        self._agent_tasks: dict[str, str] = {}
+
+    def is_available(self) -> bool:
+        """子进程后端始终可用。"""
+        return True
+
+    async def spawn(self, config: TeammateSpawnConfig) -> SpawnResult:
+        """通过任务管理器作为子进程生成新代理。"""
+        agent_id = f"{config.name}@{config.team}"
+
+        # 构建 CLI 命令
+        flags = _build_agent_cli_flags(
+            model=config.model,
+            permission_mode=None,
+        )
+        extra_env = _build_agent_env_vars()
+        env_prefix = " ".join(f"{k}={v!r}" for k, v in extra_env.items())
+
+        agent_cmd = _get_agent_command()
+        cmd_parts = [agent_cmd, "-m", "illusion"] + flags
+        command = f"{env_prefix} {' '.join(cmd_parts)}" if env_prefix else " ".join(cmd_parts)
+
+        # 创建任务
+        manager = get_task_manager()
+        try:
+            record = await manager.create_agent_task(
+                prompt=config.prompt,
+                description=f"Agent: {agent_id}",
+                cwd=config.cwd,
+                task_type="in_process_teammate",
+                model=config.model,
+                command=command,
+            )
+        except Exception as exc:
+            logger.error("[SubprocessBackend] Failed to spawn agent %s: %s", agent_id, exc)
+            return SpawnResult(
+                task_id="",
+                agent_id=agent_id,
+                backend_type=self.type,
+                success=False,
+                error=str(exc),
+            )
+
+        self._agent_tasks[agent_id] = record.id
+        logger.debug("[SubprocessBackend] Spawned agent %s as task %s", agent_id, record.id)
+        return SpawnResult(
+            task_id=record.id,
+            agent_id=agent_id,
+            backend_type=self.type,
+        )
+
+    async def send_message(self, agent_id: str, message: TeammateMessage) -> None:
+        """通过 stdin 向运行中的代理发送消息。"""
+        task_id = self._agent_tasks.get(agent_id)
+        if task_id is None:
+            raise ValueError(f"No active subprocess for agent {agent_id!r}")
+
+        payload = {
+            "text": message.text,
+            "from": message.from_agent,
+            "timestamp": message.timestamp,
+        }
+        if message.color:
+            payload["color"] = message.color
+        if message.summary:
+            payload["summary"] = message.summary
+
+        manager = get_task_manager()
+        await manager.write_to_task(task_id, json.dumps(payload))
+        logger.debug("[SubprocessBackend] Sent message to %s (task %s)", agent_id, task_id)
+
+    async def shutdown(self, agent_id: str, *, force: bool = False) -> bool:
+        """终止子进程代理。"""
+        task_id = self._agent_tasks.get(agent_id)
+        if task_id is None:
+            logger.warning("[SubprocessBackend] shutdown() called for unknown agent %s", agent_id)
+            return False
+
+        manager = get_task_manager()
+        try:
+            await manager.stop_task(task_id)
+        except ValueError as exc:
+            logger.debug("[SubprocessBackend] stop_task for %s: %s", task_id, exc)
+        finally:
+            self._agent_tasks.pop(agent_id, None)
+
+        logger.debug("[SubprocessBackend] Shut down agent %s (task %s)", agent_id, task_id)
+        return True
+
+    def get_task_id(self, agent_id: str) -> str | None:
+        """返回给定代理的任务管理器任务 ID。"""
+        return self._agent_tasks.get(agent_id)

illusion/swarm/team_helpers.py

@@ -0,0 +1,123 @@
+"""
+Swarm 团队辅助模块
+=================
+
+本模块提供 team_create / team_delete 相关的团队文件与目录管理能力。
+
+主要功能：
+    - 团队名称规范化
+    - 团队配置文件读写
+    - 团队任务目录初始化与重置
+    - 会话级团队清理注册
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shutil
+from pathlib import Path
+from typing import Any
+
+from illusion.config.paths import get_config_dir, get_data_dir
+
+# 团队负责人固定名称，与上游行为对齐
+TEAM_LEAD_NAME = "team-lead"
+
+# 当前会话中创建的团队集合（用于会话结束时兜底清理）
+_SESSION_CREATED_TEAMS: set[str] = set()
+
+
+def sanitize_name(name: str) -> str:
+    """将名称转换为适合文件系统与标识符的 slug。"""
+    sanitized = re.sub(r"[^a-zA-Z0-9]", "-", name).strip("-").lower()
+    return sanitized or "team"
+
+
+def get_team_dir(team_name: str) -> Path:
+    """返回团队目录路径（~/.illusion/teams/{team}/）。"""
+    return get_config_dir() / "teams" / sanitize_name(team_name)
+
+
+def get_team_file_path(team_name: str) -> Path:
+    """返回团队配置文件路径。"""
+    return get_team_dir(team_name) / "config.json"
+
+
+def get_team_tasks_dir(task_list_id: str) -> Path:
+    """返回团队任务目录路径（~/.illusion/data/tasks/{taskListId}/）。"""
+    return get_data_dir() / "tasks" / sanitize_name(task_list_id)
+
+
+def read_team_file(team_name: str) -> dict[str, Any] | None:
+    """读取团队配置文件。"""
+    path = get_team_file_path(team_name)
+    if not path.exists():
+        return None
+    with path.open("r", encoding="utf-8") as handle:
+        data = json.load(handle)
+    if isinstance(data, dict):
+        return data
+    return None
+
+
+def write_team_file(team_name: str, team_file: dict[str, Any]) -> None:
+    """写入团队配置文件。"""
+    path = get_team_file_path(team_name)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(team_file, handle, ensure_ascii=False, indent=2)
+
+
+def ensure_tasks_dir(task_list_id: str) -> None:
+    """确保团队任务目录存在。"""
+    get_team_tasks_dir(task_list_id).mkdir(parents=True, exist_ok=True)
+
+
+def reset_task_list(task_list_id: str) -> None:
+    """重置团队任务目录内容。"""
+    tasks_dir = get_team_tasks_dir(task_list_id)
+    tasks_dir.mkdir(parents=True, exist_ok=True)
+    for child in tasks_dir.iterdir():
+        if child.is_dir():
+            shutil.rmtree(child, ignore_errors=True)
+        else:
+            child.unlink(missing_ok=True)
+
+
+def register_team_for_session_cleanup(team_name: str) -> None:
+    """将团队登记到会话结束清理列表。"""
+    _SESSION_CREATED_TEAMS.add(team_name)
+
+
+def unregister_team_for_session_cleanup(team_name: str) -> None:
+    """从会话结束清理列表中移除团队。"""
+    _SESSION_CREATED_TEAMS.discard(team_name)
+
+
+def cleanup_team_directories(team_name: str) -> None:
+    """清理团队目录和团队任务目录。"""
+    team_file = read_team_file(team_name)
+    if team_file:
+        members = team_file.get("members", [])
+        if isinstance(members, list):
+            for member in members:
+                if not isinstance(member, dict):
+                    continue
+                worktree_path = member.get("worktreePath")
+                if isinstance(worktree_path, str) and worktree_path.strip():
+                    shutil.rmtree(Path(worktree_path), ignore_errors=True)
+
+    shutil.rmtree(get_team_dir(team_name), ignore_errors=True)
+    shutil.rmtree(get_team_tasks_dir(team_name), ignore_errors=True)
+
+
+async def cleanup_session_teams() -> None:
+    """清理当前会话创建但未显式删除的团队。"""
+    if not _SESSION_CREATED_TEAMS:
+        return
+    teams = list(_SESSION_CREATED_TEAMS)
+    for team_name in teams:
+        cleanup_team_directories(team_name)
+    _SESSION_CREATED_TEAMS.clear()
+

illusion/swarm/types.py

@@ -0,0 +1,159 @@
+"""
+Swarm 后端类型定义模块
+=====================
+
+本模块定义 Agent 派发功能使用的所有类型和协议。
+包括后端类型、代理生成配置、消息类型等。
+
+类型定义：
+    - BackendType: 支持的后端类型
+    - TeammateSpawnConfig: 代理生成配置
+    - SpawnResult: 生成结果
+    - TeammateMessage: 代理间消息
+
+协议：
+    - TeammateExecutor: 代理执行器协议
+
+使用示例：
+    >>> from illusion.swarm.types import BackendType, TeammateExecutor, SpawnResult
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal, Protocol, runtime_checkable
+
+
+# ---------------------------------------------------------------------------
+# 后端类型字面量
+# ---------------------------------------------------------------------------
+
+BackendType = Literal["subprocess", "in_process"]
+"""所有支持的后端类型。"""
+
+
+# ---------------------------------------------------------------------------
+# 代理生成配置
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TeammateSpawnConfig:
+    """生成代理的配置。"""
+
+    name: str
+    """人类可读的代理名称（例如 ``"researcher"``）。"""
+
+    team: str
+    """此代理所属的团队名称。"""
+
+    prompt: str
+    """代理的初始提示词/任务。"""
+
+    cwd: str
+    """代理的工作目录。"""
+
+    parent_session_id: str
+    """父会话 ID（用于转录关联）。"""
+
+    model: str | None = None
+    """此代理的模型覆盖。"""
+
+    system_prompt: str | None = None
+    """代理的系统提示词。"""
+
+    system_prompt_mode: Literal["default", "replace", "append"] | None = None
+    """如何应用系统提示词：替换或追加到默认。"""
+
+    color: str | None = None
+    """代理的可选 UI 颜色。"""
+
+    color_override: str | None = None
+    """明确的颜色覆盖（优先于 ``color``）。"""
+
+    permissions: list[str] = field(default_factory=list)
+    """授予此代理的工具权限。"""
+
+    plan_mode_required: bool = False
+    """此代理是否必须在实现前进入 plan 模式。"""
+
+    allow_permission_prompts: bool = False
+    """当为 False（默认）时，未列出的工具被自动拒绝。"""
+
+    worktree_path: str | None = None
+    """可选的 git worktree 路径，用于隔离的文件系统访问。"""
+
+    session_id: str | None = None
+    """明确的会话 ID（如果未提供则生成）。"""
+
+
+# ---------------------------------------------------------------------------
+# 生成结果和消息
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class SpawnResult:
+    """生成代理的结果。"""
+
+    task_id: str
+    """任务管理器中的任务 ID。"""
+
+    agent_id: str
+    """唯一代理标识符。"""
+
+    backend_type: BackendType
+    """用于生成此代理的后端。"""
+
+    success: bool = True
+    error: str | None = None
+
+
+@dataclass
+class TeammateMessage:
+    """发送给代理的消息。"""
+
+    text: str
+    from_agent: str
+    color: str | None = None
+    timestamp: str | None = None
+    summary: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# TeammateExecutor 协议
+# ---------------------------------------------------------------------------
+
+
+@runtime_checkable
+class TeammateExecutor(Protocol):
+    """代理执行后端的协议。
+
+    抽象化跨子进程和进程内后端的生成/消息/关闭操作。
+    """
+
+    type: BackendType
+
+    def is_available(self) -> bool:
+        """检查此后端在系统上是否可用。"""
+        ...
+
+    async def spawn(self, config: TeammateSpawnConfig) -> SpawnResult:
+        """使用给定配置生成新代理。"""
+        ...
+
+    async def send_message(self, agent_id: str, message: TeammateMessage) -> None:
+        """向运行中的代理发送消息。"""
+        ...
+
+    async def shutdown(self, agent_id: str, *, force: bool = False) -> bool:
+        """终止代理。
+
+        Args:
+            agent_id: 要终止的代理。
+            force: 如果为 True，立即杀死。如果为 False，尝试优雅关闭。
+
+        Returns:
+            如果代理成功终止返回 True。
+        """
+        ...