illusion-code 0.1.0__py3-none-any.whl

illusion/bridge/manager.py

@@ -0,0 +1,190 @@
+"""
+桥接会话管理器模块
+=================
+
+本模块提供桥接会话的跟踪和管理功能。
+
+主要功能：
+    - 管理生成的子会话
+    - 捕获会话输出
+    - 列出和读取会话状态
+
+类说明：
+    - BridgeSessionRecord: UI安全的会话快照
+    - BridgeSessionManager: 桥接会话管理器
+
+函数说明：
+    - get_bridge_manager: 获取单例会话管理器
+
+使用示例：
+    >>> from illusion.bridge import get_bridge_manager
+    >>> manager = get_bridge_manager()
+    >>> sessions = manager.list_sessions()
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from pathlib import Path
+
+from illusion.config.paths import get_data_dir
+from illusion.bridge.session_runner import SessionHandle, spawn_session
+
+
+@dataclass(frozen=True)
+class BridgeSessionRecord:
+    """UI安全的桥接会话快照
+    
+    Attributes:
+        session_id: 会话唯一标识符
+        command: 执行的命令
+        cwd: 工作目录
+        pid: 进程ID
+        status: 会话状态 (running/completed/failed)
+        started_at: 启动时间戳
+        output_path: 输出文件路径
+    """
+
+    session_id: str  # 会话ID
+    command: str  # 命令
+    cwd: str  # 工作目录
+    pid: int  # 进程ID
+    status: str  # 状态
+    started_at: float  # 启动时间
+    output_path: str  # 输出路径
+
+
+class BridgeSessionManager:
+    """管理桥接运行的子会话并捕获其输出
+    
+    Attributes:
+        _sessions: 会话ID到句柄的映射
+        _commands: 会话ID到命令的映射
+        _output_paths: 会话ID到输出文件路径的映射
+        _copy_tasks: 会话ID到异步复制任务的映射
+    
+    Example:
+        >>> manager = BridgeSessionManager()
+        >>> handle = await manager.spawn(session_id="test", command="echo hello", cwd=".")
+    """
+
+    def __init__(self) -> None:
+        self._sessions: dict[str, SessionHandle] = {}  # 会话映射
+        self._commands: dict[str, str] = {}  # 命令映射
+        self._output_paths: dict[str, Path] = {}  # 输出路径映射
+        self._copy_tasks: dict[str, asyncio.Task[None]] = {}  # 复制任务映射
+
+    async def spawn(self, *, session_id: str, command: str, cwd: str | Path) -> SessionHandle:
+        """生成新的桥接会话
+        
+        Args:
+            session_id: 会话唯一标识符
+            command: 要执行的命令
+            cwd: 工作目录
+        
+        Returns:
+            SessionHandle: 会话句柄
+        """
+        handle = await spawn_session(session_id=session_id, command=command, cwd=cwd)  # 生成会话
+        self._sessions[session_id] = handle  # 存储句柄
+        self._commands[session_id] = command  # 存储命令
+        output_dir = get_data_dir() / "bridge"  # 输出目录
+        output_dir.mkdir(parents=True, exist_ok=True)  # 创建目录
+        output_path = output_dir / f"{session_id}.log"  # 输出文件
+        output_path.write_text("", encoding="utf-8")  # 初始化文件
+        self._output_paths[session_id] = output_path  # 存储路径
+        self._copy_tasks[session_id] = asyncio.create_task(self._copy_output(session_id, handle))  # 启动输出复制
+        return handle
+
+    def list_sessions(self) -> list[BridgeSessionRecord]:
+        """列出所有会话
+        
+        Returns:
+            list[BridgeSessionRecord]: 按启动时间倒序排列的会话列表
+        """
+        items: list[BridgeSessionRecord] = []  # 结果列表
+        for session_id, handle in self._sessions.items():  # 遍历会话
+            process = handle.process  # 进程
+            if process.returncode is None:  # 运行中
+                status = "running"
+            elif process.returncode == 0:  # 正常退出
+                status = "completed"
+            else:  # 异常退出
+                status = "failed"
+            items.append(
+                BridgeSessionRecord(
+                    session_id=session_id,  # 会话ID
+                    command=self._commands.get(session_id, ""),  # 命令
+                    cwd=str(handle.cwd),  # 工作目录
+                    pid=process.pid or 0,  # 进程ID
+                    status=status,  # 状态
+                    started_at=handle.started_at,  # 启动时间
+                    output_path=str(self._output_paths[session_id]),  # 输出路径
+                )
+            )
+        return sorted(items, key=lambda item: item.started_at, reverse=True)  # 按时间倒序
+
+    def read_output(self, session_id: str, *, max_bytes: int = 12000) -> str:
+        """读取会话输出
+        
+        Args:
+            session_id: 会话ID
+            max_bytes: 最大返回字节数
+        
+        Returns:
+            str: 输出内容 (如果超过max_bytes，返回最后max_bytes字节)
+        """
+        path = self._output_paths.get(session_id)  # 获取输出路径
+        if path is None or not path.exists():  # 不存在
+            return ""
+        content = path.read_text(encoding="utf-8", errors="replace")  # 读取内容
+        if len(content) > max_bytes:  # 超过限制
+            return content[-max_bytes:]  # 返回最后部分
+        return content  # 返回全部
+
+    async def stop(self, session_id: str) -> None:
+        """停止会话
+        
+        Args:
+            session_id: 会话ID
+        
+        Raises:
+            ValueError: 如果会话不存在
+        """
+        handle = self._sessions.get(session_id)  # 获取句柄
+        if handle is None:  # 不存在
+            raise ValueError(f"Unknown bridge session: {session_id}")
+        await handle.kill()  # 终止会话
+
+    async def _copy_output(self, session_id: str, handle: SessionHandle) -> None:
+        """异步复制会话输出到文件
+        
+        Args:
+            session_id: 会话ID
+            handle: 会话句柄
+        """
+        path = self._output_paths[session_id]  # 输出路径
+        if handle.process.stdout is not None:  # 有输出
+            while True:
+                chunk = await handle.process.stdout.read(4096)  # 读取块
+                if not chunk:  # 无数据
+                    break
+                with path.open("ab") as stream:  # 追加模式
+                    stream.write(chunk)  # 写入
+        await handle.process.wait()  # 等待进程结束
+
+
+_DEFAULT_MANAGER: BridgeSessionManager | None = None  # 默认管理器单例
+
+
+def get_bridge_manager() -> BridgeSessionManager:
+    """获取单例桥接会话管理器
+    
+    Returns:
+        BridgeSessionManager: 全局会话管理器实例
+    """
+    global _DEFAULT_MANAGER  # 声明全局变量
+    if _DEFAULT_MANAGER is None:  # 未初始化
+        _DEFAULT_MANAGER = BridgeSessionManager()  # 创建实例
+    return _DEFAULT_MANAGER  # 返回实例

illusion/bridge/session_runner.py

@@ -0,0 +1,84 @@
+"""
+桥接会话运行器模块
+=================
+
+本模块提供桥接会话的生成和管理功能。
+
+主要功能：
+    - 生成子会话进程
+    - 管理会话生命周期
+
+类说明：
+    - SessionHandle: 生成的桥接会话句柄
+
+函数说明：
+    - spawn_session: 生成新的桥接会话
+
+使用示例：
+    >>> from illusion.bridge import spawn_session, SessionHandle
+    >>> handle = await spawn_session(session_id="test", command="echo hello", cwd=".")
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from illusion.utils.shell import create_shell_subprocess
+
+
+@dataclass
+class SessionHandle:
+    """生成的桥接会话句柄
+    
+    Attributes:
+        session_id: 会话唯一标识符
+        process: 异步子进程对象
+        cwd: 工作目录路径
+        started_at: 启动时间戳
+    """
+
+    session_id: str  # 会话ID
+    process: asyncio.subprocess.Process  # 异步进程
+    cwd: Path  # 工作目录
+    started_at: float = field(default_factory=time.time)  # 启动时间
+
+    async def kill(self) -> None:
+        """终止会话进程
+        
+        先尝试优雅终止 (terminate)，超时后强制终止 (kill)
+        """
+        self.process.terminate()  # 发送终止信号
+        try:
+            await asyncio.wait_for(self.process.wait(), timeout=3)  # 等待进程终止
+        except asyncio.TimeoutError:  # 超时
+            self.process.kill()  # 强制终止
+            await self.process.wait()  # 等待进程
+
+
+async def spawn_session(
+    *,
+    session_id: str,
+    command: str,
+    cwd: str | Path,
+) -> SessionHandle:
+    """生成一个桥接管理的子会话
+    
+    Args:
+        session_id: 会话唯一标识符
+        command: 要执行的命令
+        cwd: 工作目录
+    
+    Returns:
+        SessionHandle: 会话句柄
+    """
+    resolved_cwd = Path(cwd).resolve()  # 解析为绝对路径
+    process = await create_shell_subprocess(  # 创建子进程
+        command,  # 命令
+        cwd=resolved_cwd,  # 工作目录
+        stdout=asyncio.subprocess.PIPE,  # 标准输出管道
+        stderr=asyncio.subprocess.PIPE,  # 标准错误管道
+    )
+    return SessionHandle(session_id=session_id, process=process, cwd=resolved_cwd)

illusion/bridge/types.py

@@ -0,0 +1,113 @@
+"""
+桥接配置类型模块
+================
+
+本模块定义 Illusion Bridge 所需的配置数据类型，用于在组件之间传递配置信息。
+
+主要功能：
+    - 定义工作数据类型（WorkData）
+    - 定义工作密钥类型（WorkSecret）
+    - 定义桥接配置类型（BridgeConfig）
+
+类说明：
+    - WorkData: 工作项元数据，包含类型和ID
+    - WorkSecret: 解码后的工作密钥，包含版本、令牌和API地址
+    - BridgeConfig: 桥接配置，包含目录、机器名、会话数等
+
+使用示例：
+    >>> from illusion.bridge.types import WorkData, WorkSecret, BridgeConfig
+    >>> 
+    >>> # 创建工作数据
+    >>> work = WorkData(type="session", id="abc123")
+    >>> 
+    >>> # 创建工作密钥
+    >>> secret = WorkSecret(version=1, session_ingress_token="token", api_base_url="https://api.example.com")
+    >>> 
+    >>> # 创建桥接配置
+    >>> config = BridgeConfig(dir="/tmp/bridge", machine_name="machine-1")
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+# 默认会话超时时间：24小时（毫秒）
+DEFAULT_SESSION_TIMEOUT_MS = 24 * 60 * 60 * 1000
+
+
+@dataclass(frozen=True)
+class WorkData:
+    """
+    工作项元数据
+    
+    用于标识工作项的类型和唯一标识符。
+    
+    Attributes:
+        type: 工作类型，可选值为 "session"（会话）或 "healthcheck"（健康检查）
+        id: 工作项的唯一标识符
+    
+    使用示例：
+        >>> WorkData(type="session", id="abc123")
+    """
+
+    # 工作类型：session 或 healthcheck
+    type: Literal["session", "healthcheck"]
+    # 工作项唯一标识符
+    id: str
+
+
+@dataclass(frozen=True)
+class WorkSecret:
+    """
+    解码后的工作密钥
+    
+    包含与后端服务建立连接所需的认证信息。
+    
+    Attributes:
+        version: 密钥版本号，目前仅支持版本 1
+        session_ingress_token: 会话入口令牌，用于身份验证
+        api_base_url: API 基础地址，服务端点
+    
+    使用示例：
+        >>> WorkSecret(version=1, session_ingress_token="token", api_base_url="https://api.example.com")
+    """
+
+    # 密钥版本号
+    version: int
+    # 会话入口令牌
+    session_ingress_token: str
+    # API 基础地址
+    api_base_url: str
+
+
+@dataclass(frozen=True)
+class BridgeConfig:
+    """
+    桥接配置
+    
+    定义桥接服务的基本配置信息。
+    
+    Attributes:
+        dir: 桥接目录路径
+        machine_name: 机器名称，用于标识当前节点
+        max_sessions: 最大会话数，默认为 1
+        verbose: 是否启用详细输出，默认为 False
+        session_timeout_ms: 会话超时时间（毫秒），默认为 24 小时
+    
+    使用示例：
+        >>> config = BridgeConfig(dir="/tmp/bridge", machine_name="machine-1")
+        >>> config = BridgeConfig(dir="/tmp/bridge", machine_name="machine-1", max_sessions=5, verbose=True)
+    """
+
+    # 桥接目录路径
+    dir: str
+    # 机器名称
+    machine_name: str
+    # 最大会话数
+    max_sessions: int = 1
+    # 是否启用详细输出
+    verbose: bool = False
+    # 会话超时时间（毫秒）
+    session_timeout_ms: int = DEFAULT_SESSION_TIMEOUT_MS

illusion/bridge/work_secret.py

@@ -0,0 +1,131 @@
+"""
+工作密钥处理模块
+================
+
+本模块提供工作密钥的编码、解码和URL构建功能，用于客户端与后端服务的安全通信。
+
+主要功能：
+    - 编码工作密钥为 Base64URL 格式
+    - 解码并验证 Base64URL 格式的工作密钥
+    - 构建会话入口 WebSocket URL
+
+函数说明：
+    - encode_work_secret: 将 WorkSecret 对象编码为字符串
+    - decode_work_secret: 将字符串解码为 WorkSecret 对象
+    - build_sdk_url: 构建会话入口 WebSocket URL
+
+使用示例：
+    >>> from illusion.bridge.types import WorkSecret
+    >>> from illusion.bridge.work_secret import encode_work_secret, decode_work_secret, build_sdk_url
+    >>> 
+    >>> # 编码工作密钥
+    >>> secret = WorkSecret(version=1, session_ingress_token="token", api_base_url="https://api.example.com")
+    >>> encoded = encode_work_secret(secret)
+    >>> 
+    >>> # 解码工作密钥
+    >>> decoded = decode_work_secret(encoded)
+    >>> 
+    >>> # 构建 WebSocket URL
+    >>> ws_url = build_sdk_url("https://api.example.com", "session-123")
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+
+from illusion.bridge.types import WorkSecret
+
+
+def encode_work_secret(secret: WorkSecret) -> str:
+    """
+    编码工作密钥
+    
+    将 WorkSecret 对象序列化为 JSON 字符串，然后进行 Base64URL 编码。
+    
+    Args:
+        secret: WorkSecret 对象，包含版本、令牌和 API 地址
+    
+    Returns:
+        str: Base64URL 编码后的字符串（不包含填充字符）
+    
+    使用示例：
+        >>> secret = WorkSecret(version=1, session_ingress_token="token", api_base_url="https://api.example.com")
+        >>> encoded = encode_work_secret(secret)
+    """
+    # 将 WorkSecret 对象转换为字典并序列化为 JSON 格式
+    data = json.dumps(secret.__dict__, separators=(",", ":")).encode("utf-8")
+    # 使用 Base64URL 编码并移除填充字符
+    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")
+
+
+def decode_work_secret(secret: str) -> WorkSecret:
+    """
+    解码并验证工作密钥
+    
+    将 Base64URL 编码的字符串解码为 WorkSecret 对象，并进行版本和字段验证。
+    
+    Args:
+        secret: Base64URL 编码的工作密钥字符串
+    
+    Returns:
+        WorkSecret: 解码后的 WorkSecret 对象
+    
+    Raises:
+        ValueError: 当密钥版本不支持或缺少必需字段时
+    
+    使用示例：
+        >>> decoded = decode_work_secret("eyJ2ZXJzaW9uIjoxLCJzZXNzaW9uX2luZ3Jlc3NfdG9rZW4iOiJ0b2tlbiIsImFwaV9iYXNlX3VybCI6Imh0dHBzOi8vYXBpLmV4YW1wbGUuY29tIn0")
+    """
+    # 计算并添加缺失的填充字符
+    padding = "=" * (-len(secret) % 4)
+    # 解码 Base64URL 数据
+    raw = base64.urlsafe_b64decode((secret + padding).encode("utf-8"))
+    # 解析 JSON 数据
+    data = json.loads(raw.decode("utf-8"))
+    # 验证密钥版本
+    if data.get("version") != 1:
+        raise ValueError(f"Unsupported work secret version: {data.get('version')}")
+    # 验证会话入口令牌
+    if not data.get("session_ingress_token"):
+        raise ValueError("Invalid work secret: missing session_ingress_token")
+    # 验证 API 基础地址
+    if not isinstance(data.get("api_base_url"), str):
+        raise ValueError("Invalid work secret: missing api_base_url")
+    # 创建 WorkSecret 对象
+    return WorkSecret(
+        version=data["version"],
+        session_ingress_token=data["session_ingress_token"],
+        api_base_url=data["api_base_url"],
+    )
+
+
+def build_sdk_url(api_base_url: str, session_id: str) -> str:
+    """
+    构建会话入口 WebSocket URL
+    
+    根据 API 基础地址和会话 ID 构建完整的 WebSocket 连接 URL。
+    本地环境使用 ws 协议和 v2 版本，生产环境使用 wss 协议和 v1 版本。
+    
+    Args:
+        api_base_url: API 基础地址，如 "https://api.example.com"
+        session_id: 会话唯一标识符
+    
+    Returns:
+        str: 完整的 WebSocket URL
+    
+    使用示例：
+        >>> url = build_sdk_url("https://api.example.com", "session-123")
+        >>> # 本地: ws://api.example.com/v2/session_ingress/ws/session-123
+        >>> # 生产: wss://api.example.com/v1/session_ingress/ws/session-123
+    """
+    # 判断是否为本地环境
+    is_local = "localhost" in api_base_url or "127.0.0.1" in api_base_url
+    # 根据环境选择 WebSocket 协议
+    protocol = "ws" if is_local else "wss"
+    # 根据环境选择版本
+    version = "v2" if is_local else "v1"
+    # 提取主机地址（移除协议前缀和尾部斜杠）
+    host = api_base_url.replace("https://", "").replace("http://", "").rstrip("/")
+    # 构建完整的 WebSocket URL
+    return f"{protocol}://{host}/{version}/session_ingress/ws/{session_id}"