Moonshot AI + kvcache-ai 开源 AgentENV —— 基于 Firecracker 的智能体强化学习环境扩展平台深度解析

# 传统 RL：环境是内存中的函数
obs, reward, done, info = env.step(action)

# AgentENV 中 Jailer 配置的伪代码示意
# 实际实现使用 Rust，此处用 Python 说明逻辑

def configure_jailer(microvm_id: str, resources: ResourceSpec):
    """配置 Firecracker Jailer 的隔离参数"""
    jailer_cfg = {
        # 1. 创建独立的命名空间
        "namespaces": {
            "mount": True,      # 独立挂载树
            "pid": True,        # 独立 PID 空间
            "net": True,        # 独立网络栈
            "ipc": True,        # 独立 IPC 资源
            "uts": True,        # 独立主机名
        },
        # 2. cgroups 资源限制
        "cgroups": {
            "cpu_quota": resources.cpu_quota_us,     # 微秒级 CPU 配额
            "memory_max": resources.memory_mb,        # 最大内存（含气球）
            "memory_swap_max": 0,                      # 禁用 swap
            "io_weight": resources.io_weight,          # IO 优先级
            "pids_max": resources.max_pids,            # 最大进程数
        },
        # 3. seccomp 过滤器（系统调用白名单）
        "seccomp": {
            "default_action": "KILL",  # 白名单外的系统调用直接杀死进程
            "allowed_syscalls": [
                # 仅允许 microVM 运行所需的最小系统调用集
                "read", "write", "openat", "close",
                "mmap", "munmap", "mprotect",
                "futex", "clock_gettime", "nanosleep",
                "ioctl",   # 用于 KVM 接口
                # ... 约 50 个系统调用
            ],
        },
        # 4. chroot 到专用目录
        "chroot_dir": f"/var/lib/aenv/jailer/{microvm_id}/",
        # 5. 以非 root 用户运行
        "uid": 1000 + hash(microvm_id) % 60000,
        "gid": 1000 + hash(microvm_id) % 60000,
    }
    return jailer_cfg

# ublk + overlaybd 的 I/O 流程示意
# 实际实现使用 Rust 的 ublk 绑定

def handle_ublk_io(request: UblkRequest):
    """处理 Guest 发起的块设备 I/O 请求"""
    sector = request.sector          # 请求的扇区号
    nr_sectors = request.nr_sectors  # 请求的扇区数
    is_write = request.is_write      # 读/写

    # 1. 将扇区号映射到 overlaybd 层
    layer_info = overlaybd.lookup(sector, nr_sectors)

    if is_write:
        # 写操作：写入 COW 上层
        # 如果该扇区首次写入，先复制基础层数据到上层（COW）
        if not cow_layer.has_sector(sector):
            base_data = base_layer.read(sector, nr_sectors)
            cow_layer.write(sector, nr_sectors, base_data)
        # 写入新数据到 COW 上层
        cow_layer.write(sector, nr_sectors, request.data)
        # 标记该扇区为脏
        dirty_bitmap.mark(sector, nr_sectors)
    else:
        # 读操作：优先读 COW 上层，未命中则读基础层
        if cow_layer.has_sector(sector):
            data = cow_layer.read(sector, nr_sectors)
        else:
            data = base_layer.read(sector, nr_sectors)

    # 2. 使用 io_uring 进行异步 I/O
    # 3. 使用 Direct I/O 避免宿主机双重缓存
    return data

// AgentENV 中 ublk 队列处理的 Rust 示意代码
// 基于 io_uring 的异步 I/O 处理

use io_uring::{IoUring, opcode, types};

struct UblkQueue {
    ring: IoUring,
    sqes: Vec<opcode::ReadWrite>,
    /// 处理一批 I/O 请求
    async fn process_io_requests(&mut self, requests: Vec<UblkRequest>) {
        for (i, req) in requests.iter().enumerate() {
            let (buf, offset) = self.get_buffer_and_offset(req);
            let sqe = if req.is_write {
                opcode::Write::new(
                    types::Fixed::new(self.fd),
                    buf,
                    offset,
                ).build()
            } else {
                opcode::Read::new(
                    types::Fixed::new(self.fd),
                    buf,
                    offset,
                ).build()
            };
            // 提交到 io_uring 的提交队列
            unsafe { self.ring.submission().push(&sqe).unwrap(); }
        }
        // 等待完成
        self.ring.submit_and_wait(requests.len()).unwrap();
        // 处理完成队列
        for cqe in self.ring.completion() {
            let result = cqe.result();
            // 处理完成事件
            self.complete_io(cqe.user_data(), result);
        }
    }
}

# 增量快照的 Python 示意代码
# 实际 AgentENV 使用 Rust 调用 Firecracker API

class IncrementalSnapshotManager:
    """增量快照管理器"""

    def __init__(self, microvm_id: str):
        self.microvm_id = microvm_id
        self.base_snapshot = None   # 基础快照
        self.delta_snapshots = []   # 增量快照链
        self.last_dirty_bitmap = None  # 上次脏页位图

    async def create_snapshot(self, vm: FirecrackerMicroVM):
        """创建增量快照"""
        # 1. 暂停 VM（冻结 vCPU）
        await vm.pause()

        # 2. 获取自上次快照以来的脏页位图
        dirty_bitmap = await vm.get_dirty_log()

        if self.base_snapshot is None:
            # 首次快照：记录全量内存
            mem_snapshot = await vm.dump_memory_full()
            disk_snapshot = await vm.dump_disk_full()
            self.base_snapshot = Snapshot(
                memory=mem_snapshot,
                disk=disk_snapshot,
                bitmap=dirty_bitmap,
                timestamp=time.now(),
            )
        else:
            # 增量快照：仅记录脏页
            new_dirty_pages = self._compute_delta(
                dirty_bitmap, self.last_dirty_bitmap
            )
            mem_delta = await vm.dump_memory_regions(new_dirty_pages)
            disk_delta = await vm.dump_disk_changes(new_dirty_pages)
            delta = DeltaSnapshot(
                base_id=self.base_snapshot.id,
                dirty_pages=new_dirty_pages,
                memory_delta=mem_delta,
                disk_delta=disk_delta,
                timestamp=time.now(),
            )
            self.delta_snapshots.append(delta)

        self.last_dirty_bitmap = dirty_bitmap

        # 3. 恢复 VM
        await vm.resume()

        return self.base_snapshot or self.delta_snapshots[-1]

    async def restore_from_snapshot(self, snapshot_id: str):
        """从快照恢复环境"""
        # 1. 从基础快照加载全量状态
        vm_state = self.base_snapshot.memory.copy()
        disk_state = self.base_snapshot.disk.copy()

        # 2. 按顺序应用增量快照
        for delta in self.delta_snapshots:
            if delta.base_id == self.base_snapshot.id:
                for page_addr, page_data in delta.memory_delta.items():
                    vm_state[page_addr] = page_data
                for block_addr, block_data in delta.disk_delta.items():
                    disk_state[block_addr] = block_data

        # 3. 加载到新的 Firecracker 实例
        new_vm = await create_firecracker_vm()
        await new_vm.load_memory(vm_state)
        await new_vm.load_disk(disk_state)
        await new_vm.resume()

        return new_vm

    def _compute_delta(self, current, previous):
        """计算两个位图的差异"""
        if previous is None:
            return current
        delta = Bitmap()
        for i in range(len(current)):
            if current[i] and not previous[i]:
                delta.set(i)
        return delta

#!/usr/bin/env python3
"""AgentENV Firecracker API 操作示例 — 完整的 microVM 生命周期管理"""

import asyncio
import json
import aiohttp
from dataclasses import dataclass
from typing import Optional

@dataclass
class MicroVMConfig:
    """microVM 配置"""
    kernel_path: str = "/opt/aenv/kernel/vmlinux-6.8"
    rootfs_path: str = "/opt/aenv/images/ubuntu-22.04.ext4"
    vcpu_count: int = 2
    mem_size_mib: int = 4096
    jailer_cfg: Optional[dict] = None

class FirecrackerAPI:
    """Firecracker 控制面 API 封装"""

    def __init__(self, fc_socket_path: str):
        self.socket_path = fc_socket_path
        self.session = aiohttp.ClientSession(
            connector=aiohttp.UnixConnector(path=fc_socket_path)
        )

    async def create_vm(self, config: MicroVMConfig) -> dict:
        """创建并启动 microVM"""
        # 1. 配置内核和根文件系统
        await self._put("/boot-source", {
            "kernel_image_path": config.kernel_path,
            "boot_args": "console=ttyS0 reboot=k panic=1 pci=off"
        })

        # 2. 配置根文件系统驱动
        await self._put("/drives/rootfs", {
            "drive_id": "rootfs",
            "path_on_host": config.rootfs_path,
            "is_root_device": True,
            "is_read_only": False,
        })

        # 3. 配置 vCPU 和内存
        await self._put("/machine-config", {
            "vcpu_count": config.vcpu_count,
            "mem_size_mib": config.mem_size_mib,
            "smt": False,
            "track_dirty_pages": True,  # 启用脏页追踪（快照所需）
        })

        # 4. 配置网络接口
        await self._put("/network-interfaces/eth0", {
            "iface_id": "eth0",
            "host_dev_name": "tap0",
            "guest_mac": "02:fc:00:00:00:01",
        })

        # 5. 启动 microVM
        await self._put("/actions", {
            "action_type": "InstanceStart"
        })

        return {"status": "running", "vm_id": "vm-1"}

    async def create_snapshot(self, snapshot_id: str, mem_path: str,
                              snapshot_path: str) -> dict:
        """创建内存和磁盘快照"""
        # 先暂停 VM
        await self._put("/actions", {"action_type": "Pause"})

        # 创建快照
        await self._put("/snapshot/create", {
            "snapshot_type": "Full",
            "snapshot_path": snapshot_path,
            "mem_file_path": mem_path,
            "version": "1.0",
        })

        # 恢复 VM
        await self._put("/actions", {"action_type": "Resume"})

        return {"snapshot_id": snapshot_id, "status": "created"}

    async def restore_from_snapshot(self, mem_path: str,
                                    snapshot_path: str) -> dict:
        """从快照恢复 VM"""
        await self._put("/snapshot/load", {
            "snapshot_path": snapshot_path,
            "mem_backend": {
                "backend_type": "File",
                "backend_path": mem_path,
            },
            "resume_vm": True,
        })

        return {"status": "restored"}

    async def _put(self, path: str, body: dict) -> dict:
        """发送 PUT 请求到 Firecracker API"""
        url = f"http://localhost{path}"
        async with self.session.put(url, json=body) as resp:
            if resp.status >= 400:
                text = await resp.text()
                raise RuntimeError(f"API error {resp.status}: {text}")
            if resp.status == 204:
                return {}
            return await resp.json()

    async def close(self):
        await self.session.close()

# 使用示例
async def main():
    # 创建并启动一个 microVM
    fc_api = FirecrackerAPI("/tmp/firecracker.socket")
    vm = await fc_api.create_vm(MicroVMConfig(
        vcpu_count=4,
        mem_size_mib=8192,
    ))
    print(f"VM 创建成功: {vm}")

    # 运行一些任务...
    await asyncio.sleep(10)

    # 创建快照
    snap = await fc_api.create_snapshot(
        snapshot_id="snap-001",
        mem_path="/snapshots/snap-001/mem",
        snapshot_path="/snapshots/snap-001/vmstate",
    )
    print(f"快照创建成功: {snap}")

    # 从快照恢复到新的 VM（Fork）
    restored = await fc_api.restore_from_snapshot(
        mem_path="/snapshots/snap-001/mem",
        snapshot_path="/snapshots/snap-001/vmstate",
    )
    print(f"VM 恢复成功: {restored}")

    await fc_api.close()

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

#!/usr/bin/env python3
"""基于 AgentENV 的分布式 GRPO Agentic RL 训练示例"""

import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any
import numpy as np
import aiohttp

# ============================================================
# AgentENV 客户端（E2B 兼容 API）
# ============================================================

class AgentENVClient:
    """AgentENV 沙箱客户端（使用 E2B 兼容 API）"""

    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = aiohttp.ClientSession()

    async def create_sandbox(
        self,
        template: str = "ubuntu:22.04",
        vcpu: int = 2,
        memory_mb: int = 4096,
        timeout_s: int = 300,
    ) -> str:
        """创建新的沙箱环境"""
        async with self.session.post(
            f"{self.base_url}/sandboxes",
            json={
                "templateID": template,
                "vcpu": vcpu,
                "memoryMB": memory_mb,
                "timeout": timeout_s,
                "autoPause": True,
            }
        ) as resp:
            data = await resp.json()
            return data["sandboxID"]

    async def exec_command(
        self, sandbox_id: str, command: str
    ) -> Dict[str, Any]:
        """在沙箱中执行命令"""
        async with self.session.post(
            f"{self.base_url}/sandboxes/{sandbox_id}/execute",
            json={"command": command}
        ) as resp:
            return await resp.json()

    async def write_file(
        self, sandbox_id: str, path: str, content: str
    ) -> None:
        """写入文件到沙箱"""
        async with self.session.post(
            f"{self.base_url}/sandboxes/{sandbox_id}/files/write",
            json={"path": path, "content": content}
        ) as resp:
            resp.raise_for_status()

    async def read_file(
        self, sandbox_id: str, path: str
    ) -> str:
        """从沙箱读取文件"""
        async with self.session.get(
            f"{self.base_url}/sandboxes/{sandbox_id}/files/read",
            params={"path": path}
        ) as resp:
            data = await resp.json()
            return data["content"]

    async def pause_sandbox(self, sandbox_id: str) -> None:
        """暂停沙箱（释放 CPU/内存）"""
        async with self.session.post(
            f"{self.base_url}/sandboxes/{sandbox_id}/pause"
        ) as resp:
            resp.raise_for_status()

    async def resume_sandbox(self, sandbox_id: str) -> None:
        """恢复沙箱"""
        async with self.session.post(
            f"{self.base_url}/sandboxes/{sandbox_id}/resume"
        ) as resp:
            resp.raise_for_status()

    async def fork_sandbox(
        self, sandbox_id: str, count: int = 4
    ) -> List[str]:
        """从现有沙箱 fork 出 N 个子沙箱"""
        async with self.session.post(
            f"{self.base_url}/sandboxes/{sandbox_id}/fork",
            json={"count": count}
        ) as resp:
            data = await resp.json()
            return data["childSandboxIDs"]

    async def delete_sandbox(self, sandbox_id: str) -> None:
        """删除沙箱"""
        async with self.session.delete(
            f"{self.base_url}/sandboxes/{sandbox_id}"
        ) as resp:
            resp.raise_for_status()

    async def close(self):
        await self.session.close()


# ============================================================
# 任务定义：编码任务环境
# ============================================================

@dataclass
class CodingTask:
    """编码任务定义"""
    repo_url: str
    branch: str = "main"
    description: str = ""
    test_command: str = "pytest tests/"
    setup_commands: List[str] = field(default_factory=list)

# 示例任务：修复 Python 项目中的 bug
BENCHMARK_TASKS = [
    CodingTask(
        repo_url="https://github.com/demo/fastapi-bug-fix.git",
        branch="main",
        description="修复 FastAPI 路由中的 SQL 注入漏洞",
        test_command="pytest tests/test_security.py -v",
        setup_commands=[
            "pip install -r requirements.txt",
            "pip install pytest",
        ],
    ),
    CodingTask(
        repo_url="https://github.com/demo/react-form-bug.git",
        branch="main",
        description="修复 React 表单验证组件中的状态更新 bug",
        test_command="npm test",
        setup_commands=[
            "npm install",
        ],
    ),
]


# ============================================================
# GRPO 训练器
# ============================================================

@dataclass
class GRPOConfig:
    """GRPO 训练配置"""
    group_size: int = 8          # 每组采样数（GRPO 的 G）
    num_iterations: int = 100    # 训练迭代次数
    learning_rate: float = 1e-5
    kl_coef: float = 0.01
    clip_epsilon: float = 0.2
    max_steps_per_task: int = 20
    fork_count: int = 8          # 每次 fork 子沙箱数

class GRPOTrainer:
    """GRPO 算法训练器——使用 AgentENV 作为环境后端"""

    def __init__(
        self,
        policy_model: Any,  # 实际使用会传入 HuggingFace 模型
        tokenizer: Any,
        config: GRPOConfig,
        env_client: AgentENVClient,
    ):
        self.policy = policy_model
        self.tokenizer = tokenizer
        self.config = config
        self.env = env_client
        self.optimizer = torch.optim.AdamW(
            self.policy.parameters(),
            lr=config.learning_rate,
        )

    async def setup_environment(self, task: CodingTask) -> str:
        """构建训练环境并返回沙箱 ID"""
        # 1. 创建基础沙箱
        sandbox_id = await self.env.create_sandbox(
            template="coding-agent:latest",
            vcpu=4,
            memory_mb=8192,
            timeout_s=3600,
        )

        # 2. 克隆代码仓库
        await self.env.exec_command(
            sandbox_id,
            f"git clone --branch {task.branch} {task.repo_url} /workspace"
        )

        # 3. 执行安装命令
        for cmd in task.setup_commands:
            result = await self.env.exec_command(sandbox_id, cmd)
            if result["exitCode"] != 0:
                raise RuntimeError(
                    f"Setup failed: {result['stderr']}"
                )

        return sandbox_id

    async def collect_trajectories(
        self, base_sandbox_id: str, task: CodingTask
    ) -> List[Dict]:
        """使用 Fork 并行收集多条轨迹"""
        # 1. 从基础环境 fork 出多个子沙箱
        child_ids = await self.env.fork_sandbox(
            base_sandbox_id,
            count=self.config.fork_count,
        )
        print(f"Forked {len(child_ids)} child sandboxes from {base_sandbox_id}")

        trajectories = []
        tasks = []

        for i, child_id in enumerate(child_ids):
            tasks.append(
                self._rollout_single(child_id, task, i)
            )

        # 并行执行所有 rollout
        results = await asyncio.gather(*tasks)
        trajectories.extend(results)

        # 清理子沙箱
        for child_id in child_ids:
            await self.env.delete_sandbox(child_id)

        return trajectories

    async def _rollout_single(
        self, sandbox_id: str, task: CodingTask, rank: int
    ) -> Dict:
        """单条轨迹采样"""
        trajectory = {
            "sandbox_id": sandbox_id,
            "rank": rank,
            "steps": [],
            "final_reward": 0.0,
            "success": False,
        }

        current_state = await self._get_env_state(sandbox_id)

        for step in range(self.config.max_steps_per_task):
            # 策略模型生成动作
            action = await self._policy_generate(current_state)

            # 在沙箱中执行动作
            result = await self.env.exec_command(sandbox_id, action)

            # 获取新状态和奖励
            next_state = await self._get_env_state(sandbox_id)
            reward = self._compute_reward(result, task)

            trajectory["steps"].append({
                "state": current_state,
                "action": action,
                "result": result,
                "reward": reward,
                "next_state": next_state,
            })

            current_state = next_state

            # 检查是否完成任务
            if result["exitCode"] == 0 and "ALL TESTS PASSED" in result["stdout"]:
                trajectory["success"] = True
                trajectory["final_reward"] = 10.0
                break

        return trajectory

    async def _get_env_state(self, sandbox_id: str) -> str:
        """获取环境状态摘要"""
        # 读取当前工作目录状态
        result = await self.env.exec_command(
            sandbox_id,
            "echo '=== FILES ===' && find . -name '*.py' -newer /tmp/start | head -20 && "
            "echo '=== GIT DIFF ===' && git diff --stat 2>/dev/null || true"
        )
        return result["stdout"]

    async def _policy_generate(self, state: str) -> str:
        """策略模型根据状态生成动作（代码修改）"""
        prompt = (
            f"当前代码库状态:\n{state}\n\n"
            f"请生成下一步需要执行的 shell 命令来修复 bug。"
            f"只输出命令，不要解释。"
        )
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.policy.generate(
                **inputs,
                max_new_tokens=200,
                temperature=1.0,
                top_p=0.9,
            )
        action = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 提取纯命令部分
        action = action.split("\n")[-1].strip()
        return action

    def _compute_reward(self, result: Dict, task: CodingTask) -> float:
        """计算奖励"""
        reward = 0.0
        if result["exitCode"] == 0:
            reward += 1.0  # 命令执行成功
        if "ALL TESTS PASSED" in result.get("stdout", ""):
            reward += 5.0  # 测试通过
        if "ERROR" in result.get("stderr", ""):
            reward -= 0.5  # 出现错误
        return reward

    async def train_step(self, trajectories: List[Dict]) -> Dict:
        """GRPO 训练步骤"""
        # 1. 收集所有轨迹的奖励
        all_rewards = []
        all_log_probs = []
        all_states = []
        all_actions = []

        for traj in trajectories:
            for step in traj["steps"]:
                all_rewards.append(step["reward"])
                all_states.append(step["state"])
                all_actions.append(step["action"])

        if not all_rewards:
            return {"loss": 0.0, "mean_reward": 0.0}

        # 2. GRPO: 组内归一化奖励
        rewards = np.array(all_rewards, dtype=np.float32)
        mean_reward = rewards.mean()
        std_reward = rewards.std() + 1e-8
        normalized_rewards = (rewards - mean_reward) / std_reward

        # 3. 计算策略损失
        total_loss = 0.0
        for i in range(0, len(all_states), self.config.group_size):
            batch_end = min(i + self.config.group_size, len(all_states))
            batch_rewards = normalized_rewards[i:batch_end]

            # 重新计算当前策略的 log prob
            # 实际代码会使用 policy.forward() 计算
            # 这里简化示意
            advantages = torch.tensor(batch_rewards, dtype=torch.float32)
            surrogate_loss = -advantages.mean()
            total_loss += surrogate_loss.item()

        # 4. 反向传播
        loss_tensor = torch.tensor(total_loss, requires_grad=True)
        self.optimizer.zero_grad()
        loss_tensor.backward()
        torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 1.0)
        self.optimizer.step()

        return {
            "loss": total_loss,
            "mean_reward": float(mean_reward),
            "num_samples": len(all_rewards),
        }

    async def train(self, tasks: List[CodingTask]):
        """完整训练循环"""
        for iteration in range(self.config.num_iterations):
            print(f"\n=== Iteration {iteration + 1}/{self.config.num_iterations} ===")

            all_trajectories = []

            for task in tasks:
                # 1. 设置环境
                base_sandbox = await self.setup_environment(task)
                print(f"  Task: {task.description}")

                # 2. Fork 并并行收集轨迹
                trajectories = await self.collect_trajectories(
                    base_sandbox, task
                )
                all_trajectories.extend(trajectories)

                # 3. 清理基础环境
                await self.env.pause_sandbox(base_sandbox)

                # 4. 评估结果
                success_rate = sum(
                    1 for t in trajectories if t["success"]
                ) / len(trajectories)
                print(f"    Success rate: {success_rate:.1%}")

            # 5. GRPO 更新
            metrics = await self.train_step(all_trajectories)
            print(f"  Train metrics: loss={metrics['loss']:.4f}, "
                  f"reward={metrics['mean_reward']:.2f}")

            # 6. 保存检查点
            if (iteration + 1) % 10 == 0:
                torch.save(
                    self.policy.state_dict(),
                    f"checkpoint_iter_{iteration + 1}.pt"
                )


# ============================================================
# 主程序入口
# ============================================================

async def main():
    # 初始化 AgentENV 客户端
    client = AgentENVClient("http://localhost:8000")

    # 初始化策略模型（示意：使用小模型）
    # 实际使用会加载 HuggingFace 模型
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model_name = "Qwen/Qwen2.5-1.5B-Instruct"
    policy = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 配置训练参数
    config = GRPOConfig(
        group_size=8,
        num_iterations=100,
        learning_rate=1e-5,
        fork_count=8,
        max_steps_per_task=20,
    )

    # 创建训练器
    trainer = GRPOTrainer(
        policy_model=policy,
        tokenizer=tokenizer,
        config=config,
        env_client=client,
    )

    # 开始训练
    try:
        await trainer.train(BENCHMARK_TASKS)
    finally:
        await client.close()


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

#!/usr/bin/env python3
"""AgentENV 网络隔离配置示例"""

import ipaddress
from typing import List, Optional

class NetworkIsolationConfig:
    """沙箱网络隔离配置"""

    # 默认禁止的内部 CIDR 列表
    DEFAULT_DENIED_CIDRS = [
        "10.0.0.0/8",       # 私有网络
        "172.16.0.0/12",    # 私有网络
        "192.168.0.0/16",   # 私有网络
        "169.254.0.0/16",   # 链路本地
        "127.0.0.0/8",      # 本地回环
    ]

    def __init__(
        self,
        allow_internet: bool = True,
        denied_cidrs: Optional[List[str]] = None,
        allowed_cidrs: Optional[List[str]] = None,
    ):
        self.allow_internet = allow_internet
        self.denied_cidrs = [
            ipaddress.ip_network(cidr)
            for cidr in (denied_cidrs or self.DEFAULT_DENIED_CIDRS)
        ]
        self.allowed_cidrs = [
            ipaddress.ip_network(cidr)
            for cidr in (allowed_cidrs or [])
        ]

    def validate_egress(self, target_ip: str, target_port: int) -> bool:
        """验证出站目标是否允许"""
        ip = ipaddress.ip_address(target_ip)

        # 白名单优先：如果在 allowed_cidrs 中，允许
        for network in self.allowed_cidrs:
            if ip in network:
                return True

        # 黑名单：如果在 denied_cidrs 中，拒绝
        for network in self.denied_cidrs:
            if ip in network:
                return False

        # 默认策略：根据 allow_internet 决定
        return self.allow_internet

    def to_api_request(self) -> dict:
        """转换为 AgentENV API 请求格式"""
        return {
            "allowInternetAccess": self.allow_internet,
            "alwaysDeniedCIDRs": [str(c) for c in self.denied_cidrs],
            "alwaysAllowedCIDRs": [str(c) for c in self.allowed_cidrs],
        }


# 为不同安全级别的训练任务配置网络
def create_network_config(task_type: str) -> NetworkIsolationConfig:
    """根据任务类型创建网络隔离配置"""
    configs = {
        "high_security": NetworkIsolationConfig(
            allow_internet=False,
            denied_cidrs=[
                "0.0.0.0/0",  # 完全禁止所有网络访问
            ],
        ),
        "coding_agent": NetworkIsolationConfig(
            allow_internet=True,
            denied_cidrs=[
                "10.0.0.0/8",
                "172.16.0.0/12",
                "192.168.0.0/16",
                "100.64.0.0/10",  # 避免访问内部服务
            ],
            allowed_cidrs=[
                "0.0.0.0/0",  # 允许访问外网（如 pip install）
            ],
        ),
        "web_agent": NetworkIsolationConfig(
            allow_internet=True,
            denied_cidrs=[
                "10.0.0.0/8",
                "172.16.0.0/12",
                "192.168.0.0/16",
            ],
            allowed_cidrs=[
                "0.0.0.0/0",  # 允许访问外网
            ],
        ),
    }
    return configs.get(task_type, configs["coding_agent"])

# 方式 1: 一键安装脚本（单节点，推荐入门）
curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh | sudo bash
sudo systemctl start aenv

# 方式 2: Docker 部署
curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/docker-setup.sh | sudo bash
docker pull ghcr.io/kvcache-ai/aenv-server:latest
docker run -d --privileged -v /dev:/dev -p 8000:8000 ghcr.io/kvcache-ai/aenv-server:latest

# 方式 3: Docker Compose（模拟多节点集群）
# 见 docker-compose.yml

# 方式 4: Kubernetes（生产级多节点）
# 包含 Gateway、Scheduler、Node DaemonSet

# 方式 5: 从源码构建
git clone https://github.com/kvcache-ai/AgentENV.git
cd AgentENV
cargo build --release

# 安装 CLI
curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install-cli.sh | bash

# 认证
aenv auth
# AENV server URL [http://localhost:8000]: http://127.0.0.1:8000
# API key: dummy

# 拉取模板
aenv pull ubuntu:22.04 --name ubuntu
aenv pull docker.io/library/python:3.12 --name python312

# 启动沙箱并进入交互式 shell
aenv start ubuntu

# 启动沙箱（后台模式）
aenv start ubuntu --detach

# 列出所有沙箱
aenv ls

# 执行单条命令
aenv exec <sandbox-id> ls -la /

# 暂停/恢复
aenv pause <sandbox-id>
aenv resume <sandbox-id>

# 设置超时
aenv timeout <sandbox-id> 600

# 删除沙箱
aenv delete <sandbox-id>

# docker-compose.yml —— 多节点 AgentENV 集群
version: "3.8"

services:
  gateway:
    image: ghcr.io/kvcache-ai/aenv-gateway:latest
    ports:
      - "8080:8080"
    environment:
      - SCHEDULER_URL=http://scheduler:9090
    depends_on:
      - scheduler

  scheduler:
    image: ghcr.io/kvcache-ai/aenv-scheduler:latest
    ports:
      - "9090:9090"
    environment:
      - NODE_LABELS=region=us-east-1,zone=a

  node-1:
    image: ghcr.io/kvcache-ai/aenv-server:latest
    privileged: true
    devices:
      - /dev/kvm:/dev/kvm
    volumes:
      - /dev:/dev
      - aenv-data-1:/var/lib/aenv
    environment:
      - SCHEDULER_URL=http://scheduler:9090
      - NODE_ID=node-1
    depends_on:
      - scheduler

  node-2:
    image: ghcr.io/kvcache-ai/aenv-server:latest
    privileged: true
    devices:
      - /dev/kvm:/dev/kvm
    volumes:
      - /dev:/dev
      - aenv-data-2:/var/lib/aenv
    environment:
      - SCHEDULER_URL=http://scheduler:9090
      - NODE_ID=node-2
    depends_on:
      - scheduler

volumes:
  aenv-data-1:
  aenv-data-2:

#!/usr/bin/env python3
"""AgentENV 性能基准测试脚本"""

import asyncio
import time
import statistics
from dataclasses import dataclass, field
from typing import List
from agentenv_client import AgentENVClient  # 假设的客户端库

@dataclass
class BenchmarkResult:
    """基准测试结果"""
    cold_start_ms: List[float] = field(default_factory=list)
    snapshot_boot_ms: List[float] = field(default_factory=list)
    pause_ms: List[float] = field(default_factory=list)
    resume_ms: List[float] = field(default_factory=list)
    fork_ms: List[float] = field(default_factory=list)
    snapshot_create_ms: List[float] = field(default_factory=list)

    def summary(self) -> str:
        """生成测试摘要"""
        lines = []
        lines.append("=== AgentENV 性能基准测试 ===")
        lines.append("")
        for name, data in [
            ("冷启动时间", self.cold_start_ms),
            ("快照启动时间", self.snapshot_boot_ms),
            ("暂停时间", self.pause_ms),
            ("恢复时间", self.resume_ms),
            ("Fork 时间", self.fork_ms),
            ("快照创建时间", self.snapshot_create_ms),
        ]:
            if data:
                lines.append(
                    f"{name}: "
                    f"avg={statistics.mean(data):.1f}ms, "
                    f"p50={statistics.median(data):.1f}ms, "
                    f"p99={sorted(data)[int(len(data)*0.99)]:.1f}ms, "
                    f"min={min(data):.1f}ms, max={max(data):.1f}ms"
                )
        return "\n".join(lines)

async def run_benchmark(
    client: AgentENVClient,
    iterations: int = 50,
) -> BenchmarkResult:
    """运行基准测试"""
    result = BenchmarkResult()

    # 1. 测试冷启动
    print(f"测试冷启动 ({iterations} 次)...")
    for i in range(iterations):
        start = time.perf_counter()
        sid = await client.create_sandbox(
            template="ubuntu:22.04",
            vcpu=2,
            memory_mb=1024,
        )
        elapsed = (time.perf_counter() - start) * 1000
        result.cold_start_ms.append(elapsed)
        await client.delete_sandbox(sid)

    # 2. 先创建一个模板，然后测试快照启动
    print("创建模板快照...")
    builder = await client.create_sandbox(
        template="ubuntu:22.04",
        vcpu=2,
        memory_mb=1024,
    )
    await client.exec_command(builder, "apt-get update && apt-get install -y python3-pip")
    template_id = await client.create_template(builder, "benchmark-template")
    await client.delete_sandbox(builder)

    # 3. 测试快照启动
    print(f"测试快照启动 ({iterations} 次)...")
    for i in range(iterations):
        start = time.perf_counter()
        sid = await client.create_sandbox(template=template_id)
        elapsed = (time.perf_counter() - start) * 1000
        result.snapshot_boot_ms.append(elapsed)

        # 测试暂停
        start = time.perf_counter()
        await client.pause_sandbox(sid)
        elapsed = (time.perf_counter() - start) * 1000
        result.pause_ms.append(elapsed)

        # 测试恢复
        start = time.perf_counter()
        await client.resume_sandbox(sid)
        elapsed = (time.perf_counter() - start) * 1000
        result.resume_ms.append(elapsed)

        # 测试快照创建
        start = time.perf_counter()
        await client.create_snapshot(sid, f"snap-{i}")
        elapsed = (time.perf_counter() - start) * 1000
        result.snapshot_create_ms.append(elapsed)

        await client.delete_sandbox(sid)

    # 4. 测试 Fork
    print(f"测试 Fork ({min(iterations, 16)} 次)...")
    parent = await client.create_sandbox(template=template_id)
    start = time.perf_counter()
    children = await client.fork_sandbox(parent, count=16)
    elapsed = (time.perf_counter() - start) * 1000
    result.fork_ms.append(elapsed)
    for child in children:
        await client.delete_sandbox(child)
    await client.delete_sandbox(parent)

    return result

async def main():
    client = AgentENVClient("http://localhost:8000")
    try:
        result = await run_benchmark(client, iterations=30)
        print(result.summary())
    finally:
        await client.close()

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

# E2B SDK 代码 —— 无需修改即可在 AgentENV 上运行
import os
from e2b import Sandbox

# 只需修改这一行：
os.environ["E2B_API_URL"] = "http://localhost:8000"

# 以下代码完全不变
sandbox = Sandbox(template="ubuntu:22.04")
sandbox.commands.run("pip install requests")
code = """
import requests
response = requests.get('https://api.example.com/data')
print(response.json())
"""
sandbox.files.write("/tmp/script.py", code)
result = sandbox.commands.run("python /tmp/script.py")
print(result.stdout)
sandbox.close()