Files
codex-app-server/cloud/tests/test_docker_sandbox.py
T
host d51c06a900 Worker concurrency (capacity enforced), survival (heartbeat/sweep/drain), retry endpoint
1. Concurrency: capacity is now a real limit, not display-only.
   run_all with capacity>1 runs a manual thread pool (threading,
   not concurrent.futures: a test in this repo puts cloud/ on
   sys.path, shadowing stdlib queue and breaking
   concurrent.futures' SimpleQueue import). Claim stays CAS, so
   two threads racing one task still yield one owner; max_tasks
   counts dispatched+running to avoid overshoot; stop signals
   stop new claims only, running tasks are drained by run_loop.
   Running state moved from single fields to a _running dict
   (proc/container per task) so cancel kills the right task.
   Docker timeout is now a per-call argument, never a mutated
   shared executor field. CostStore opens a fresh connection per
   record/report (shared sqlite connection is thread-bound and
   raised ProgrammingError under the pool; :memory: keeps the
   shared connection because a fresh one is an empty DB).
   Final status write is CAS on status='running' so a background
   thread cannot overwrite an interrupted task with completed.
2. Survival: run_loop starts a daemon heartbeat thread (every
   30s, refreshes running tasks' last_heartbeat) and sweeps stale
   tasks every round (threshold 300s, best-effort); live tasks
   never go stale now, dead workers' tasks are reclaimed without
   manual /admin/sweep. On exit, run_loop waits up to
   drain_timeout (CLI --drain-timeout, default 60s, inside
   systemd's 90s stop window) for running tasks before killing
   leftovers and marking interrupted.
3. Retry: POST /tasks/{id}/retry (maintainer+, 404/403/409):
   failed/interrupted/merge-failed go back to queued with
   result/error/model cleared, prompt/model/deps kept, CAS
   update, audit task_retried.
Tests: 5 concurrency tests (barrier overlap, serial capacity=1,
max_tasks, wait_for_running, CAS final write), 3 retry tests;
full suites pass.
2026-10-26 06:46:20 +08:00

339 lines
14 KiB
Python

"""Docker 沙箱执行器单测:命令构造、安全挂载、超时与取消。
运行(在 ~/workspace/codex-app-server 下):
python3 -m unittest discover -s cloud/tests -p "test_docker_sandbox.py" -v
说明:全部用注入的假 docker runner,不调用真实 docker、不联网、
不跑真实模型回合;env 映射里的值只是占位假名,绝不是真实密钥。
"""
import json
import os
import sqlite3
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
try:
from cloud.sandbox import DockerExecutor
from cloud.server import init_db
from cloud.worker import Worker
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from cloud.sandbox import DockerExecutor
from cloud.server import init_db
from cloud.worker import Worker
def _ok_runner_factory(seen):
def fake(cmd, timeout=None):
seen.append(list(cmd))
return SimpleNamespace(
returncode=0,
stdout=json.dumps({"status": "completed", "text": "done"}),
stderr="")
return fake
def _insert_task(db_path, task_id, prompt="do work"):
now = time.time()
conn = sqlite3.connect(db_path)
try:
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" idempotency_key, depends_on, priority, created_by,"
" created_at, updated_at, last_heartbeat)"
" VALUES(?, 'p1', ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
(task_id, prompt, now, now),
)
conn.commit()
finally:
conn.close()
def _task_status(db_path, task_id):
conn = sqlite3.connect(db_path)
try:
row = conn.execute(
"SELECT status, error_text FROM tasks WHERE id = ?",
(task_id,)).fetchone()
finally:
conn.close()
return row
class DockerCommandTestCase(unittest.TestCase):
def _executor(self, seen, **kwargs):
kwargs.setdefault("runner", _ok_runner_factory(seen))
kwargs.setdefault("env", {"XCM_API_KEY": "placeholder-not-a-real-key"})
return DockerExecutor("/fake/host.py", **kwargs)
def test_command_mounts_only_task_worktree_and_ro_code(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
cmd = ex.build_run_command(
{"id": "t1", "worktree": "/tmp/wt-t1"},
["python3", "/app/host.py", "status", "--task", "t1"])
mounts = [cmd[i + 1] for i, a in enumerate(cmd) if a == "-v"]
self.assertEqual(mounts, ["/tmp/wt-t1:/work", "/fake:/app:ro"])
# 绝不挂载宿主根目录。
self.assertNotIn("/:/work", mounts)
self.assertNotIn("/:/app:ro", mounts)
self.assertIn("--rm", cmd)
self.assertIn("/tmp/wt-t1:/work", cmd)
def test_command_resource_limits_and_non_root_user(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen, cpus="1", memory="512m",
user="1000:1000")
cmd = ex.build_run_command(
{"id": "t1", "worktree": "/tmp/wt"},
["echo", "hi"])
self.assertEqual(cmd[cmd.index("--cpus") + 1], "1")
self.assertEqual(cmd[cmd.index("--memory") + 1], "512m")
self.assertEqual(cmd[cmd.index("--user") + 1], "1000:1000")
self.assertNotEqual(cmd[cmd.index("--user") + 1], "0")
self.assertNotEqual(cmd[cmd.index("--user") + 1], "root")
def test_command_defaults_are_2cpu_2g(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
cmd = ex.build_run_command(
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
self.assertEqual(cmd[cmd.index("--cpus") + 1], "2")
self.assertEqual(cmd[cmd.index("--memory") + 1], "2g")
def test_container_name_sanitized(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
self.assertEqual(ex.container_name({"id": "a b/c"}),
"codex-a-b-c")
cmd = ex.build_run_command(
{"id": "a b/c", "worktree": "/tmp/wt"}, ["echo", "hi"])
self.assertIn("codex-a-b-c", cmd)
self.assertEqual(cmd[cmd.index("--name") + 1], "codex-a-b-c")
def test_env_name_only_no_value_in_command(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
cmd = ex.build_run_command(
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
joined = " ".join(cmd)
self.assertIn("-e", cmd)
self.assertIn("XCM_API_KEY", cmd)
# 值绝不出现在命令列表里(-e 只传名,docker 从宿主环境取值)。
self.assertNotIn("placeholder-not-a-real-key", joined)
self.assertNotIn("placeholder-not-a-real-key=", joined)
def test_missing_key_does_not_start_container(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = DockerExecutor("/fake/host.py",
runner=_ok_runner_factory(seen), env={})
with self.assertRaises(RuntimeError) as ctx:
ex.dispatch_task({"id": "t1", "prompt": "hi",
"worktree": "/tmp/wt"})
self.assertIn("Missing environment variable", str(ctx.exception))
self.assertEqual(seen, [])
def test_missing_docker_does_not_call_runner(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which", return_value=None):
ex = self._executor(seen)
with self.assertRaises(RuntimeError) as ctx:
ex.dispatch_task({"id": "t1", "prompt": "hi",
"worktree": "/tmp/wt"})
self.assertIn("docker not available", str(ctx.exception))
self.assertEqual(seen, [])
def test_dispatch_and_run_shapes(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
dispatched = ex.dispatch_task(
{"id": "t1", "prompt": "do it", "worktree": "/tmp/wt"})
ran = ex.run_task(
{"id": "t1", "prompt": "again", "worktree": "/tmp/wt"})
self.assertEqual(dispatched["status"], "completed")
self.assertEqual(dispatched["executor"], "docker")
dispatch_cmd = seen[0]
self.assertIn("dispatch", dispatch_cmd)
self.assertIn("t1", dispatch_cmd)
self.assertIn("do it", dispatch_cmd)
run_cmd = seen[1]
self.assertIn("run", run_cmd)
self.assertIn("again", run_cmd)
def test_execute_is_readonly_status_not_dispatch(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
result = ex.execute({"id": "t1", "worktree": "/tmp/wt"})
self.assertEqual(result["status"], "completed")
self.assertIn("status", seen[0])
self.assertNotIn("dispatch", seen[0])
# run 在命令里会出现 "--name" 等,但内层命令不应是 run:
# 内层 host.py 命令的第一个动词是 status。
inner = seen[0][seen[0].index("codex-app-server:latest") + 1:]
self.assertEqual(inner[2], "status")
def test_timeout_kills_and_removes_container(self):
seen = []
def fake(cmd, timeout=None):
seen.append(list(cmd))
if len(cmd) > 1 and cmd[1] == "run":
raise subprocess.TimeoutExpired(cmd, timeout)
return SimpleNamespace(returncode=0, stdout="", stderr="")
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = DockerExecutor("/fake/host.py", runner=fake,
env={"XCM_API_KEY": "placeholder"},
timeout=5)
result = ex.dispatch_task(
{"id": "t9", "prompt": "slow", "worktree": "/tmp/wt"})
self.assertEqual(result["status"], "failed")
self.assertIn("timeout", result["stdout"])
kinds = [(c[1], c[2]) for c in seen if len(c) > 2]
self.assertIn(("kill", "codex-t9"), kinds)
self.assertIn(("rm", "-f"), [(c[1], c[2]) for c in seen
if len(c) > 2])
def test_kill_container_calls_kill_and_rm(self):
seen = []
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = self._executor(seen)
ok = ex.kill_container({"id": "t2"})
self.assertTrue(ok)
self.assertEqual([c[1] for c in seen], ["kill", "rm"])
self.assertEqual(seen[0][2], "codex-t2")
self.assertEqual(seen[1][2], "-f")
self.assertEqual(seen[1][3], "codex-t2")
class WorkerDockerTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
self.env = {"XCM_API_KEY": "placeholder-not-a-real-key"}
def tearDown(self):
self.tmp.cleanup()
def test_worker_docker_missing_fails_task_with_docker_error(self):
seen = []
def fake(cmd, timeout=None):
seen.append(list(cmd))
return SimpleNamespace(returncode=0, stdout="{}", stderr="")
_insert_task(self.db_path, "d1", prompt="do it")
with mock.patch("cloud.sandbox.shutil.which", return_value=None):
ex = DockerExecutor("/fake/host.py", runner=fake, env=self.env)
worker = Worker(self.db_path, "/fake/host.py",
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
result = worker.run_one()
self.assertTrue(result["dispatched"])
self.assertEqual(result["status"], "failed")
self.assertIn("docker", result["error"])
self.assertEqual(seen, [])
status, error_text = _task_status(self.db_path, "d1")
self.assertEqual(status, "failed")
self.assertIn("docker", error_text)
def test_worker_docker_success(self):
seen = []
_insert_task(self.db_path, "d2", prompt="do it")
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = DockerExecutor("/fake/host.py",
runner=_ok_runner_factory(seen),
env=self.env)
worker = Worker(self.db_path, "/fake/host.py",
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
result = worker.run_one()
self.assertTrue(result["dispatched"])
self.assertEqual(result["status"], "completed")
self.assertEqual(result["executor"], "docker")
self.assertTrue(any(c[1] == "run" for c in seen))
def test_worker_cancel_kills_docker_container(self):
seen = []
_insert_task(self.db_path, "d3", prompt="do it")
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = DockerExecutor("/fake/host.py",
runner=_ok_runner_factory(seen),
env=self.env)
worker = Worker(self.db_path, "/fake/host.py",
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
# 模拟 dispatch 阻塞期间:worker 已把容器名记下来。
worker._running_set("d3", container="codex-d3")
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"UPDATE tasks SET status = 'running' WHERE id = 'd3'")
conn.commit()
finally:
conn.close()
self.assertTrue(worker.cancel("d3"))
self.assertEqual([c[1] for c in seen], ["kill", "rm"])
status, _ = _task_status(self.db_path, "d3")
self.assertEqual(status, "interrupted")
if __name__ == "__main__":
unittest.main()
class TestDockerExtraMounts(DockerCommandTestCase):
def test_extra_mounts_appended_after_base_mounts(self):
with tempfile.TemporaryDirectory() as td:
work = os.path.join(td, "work")
os.makedirs(work)
state = os.path.join(td, "host-state.json")
open(state, "w").write("{}")
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
ex = DockerExecutor("/fake/host.py", runner=_ok_runner_factory([]),
env={"XCM_API_KEY": "placeholder-not-a-real-key"},
extra_mounts=[f"{work}:/app/work", f"{state}:/app/host-state.json"])
cmd = ex.build_run_command({"id": "t1"}, ["python3", "/app/host.py", "status"])
# 基础挂载在前,额外挂载在后,且 /app/work 覆盖只读子路径
self.assertIn(f"{work}:/app/work:rw", cmd)
self.assertIn(f"{state}:/app/host-state.json:rw", cmd)
mounts = [cmd[i + 1] for i, a in enumerate(cmd) if a == "-v"]
self.assertEqual(mounts[1], "/fake:/app:ro")
self.assertIn(f"{work}:/app/work:rw", mounts)
self.assertLess(mounts.index("/fake:/app:ro"), mounts.index(f"{work}:/app/work:rw"))
def test_extra_mounts_reject_dangerous_paths(self):
with self.assertRaises(ValueError):
DockerExecutor("/fake/host.py", extra_mounts=["/: /x"])
with self.assertRaises(ValueError):
DockerExecutor("/fake/host.py", extra_mounts=["/etc/codex-app-server/env:/x"])
with self.assertRaises(ValueError):
DockerExecutor("/fake/host.py", extra_mounts=["/tmp:/app"])
with self.assertRaises(ValueError):
DockerExecutor("/fake/host.py", extra_mounts=["/nonexistent-xyz:/x"])