mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
数据目录里的文件一律先写临时文件再改名:新增 core/atomic(同目录、带随机后缀的临时文件,写完 os.replace;只该生成一次的用 os.link 建,已有就不动)
- 工件(内容寻址):以前共用 <id>.tmp,两个进程同时写同一份内容时后到的会截断先到的半成品、先到的改名又让后到的改名失败;现在各写各的临时文件 - 加密密钥 .secret_key、断言签名私钥:几个进程同时第一次启动时只建出一把,后到的读先到的那一把(以前各生成一把,各自加密、签名的东西对方认不得),文件从建出来起就是 0600 - 附件抽出文字的缓存:临时文件名带随机后缀 - 文件工具(file_write)的覆盖和追加、三种系统沙箱写进工作目录的代码和输入文件、seatbelt 的策略文件:先写临时文件再改名 - 上传表格的版本库、原件存档、记录表本来就是这样写的,不动
This commit is contained in:
@@ -23,6 +23,7 @@ from typing import Any
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings
|
||||
|
||||
#: 断言放在这个请求头里
|
||||
@@ -35,14 +36,13 @@ _KEY_FILE = ".assertion_ed25519.pem"
|
||||
@lru_cache(maxsize=1)
|
||||
def _private_key() -> Ed25519PrivateKey:
|
||||
path = settings.data_dir / _KEY_FILE
|
||||
if path.exists():
|
||||
key = serialization.load_pem_private_key(path.read_bytes(), password=None)
|
||||
assert isinstance(key, Ed25519PrivateKey)
|
||||
return key
|
||||
key = Ed25519PrivateKey.generate()
|
||||
path.write_bytes(key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8,
|
||||
serialization.NoEncryption()))
|
||||
path.chmod(0o600)
|
||||
if not path.exists():
|
||||
# 只建一次(core/atomic.create_bytes):几个进程同时第一次签名,后到的读先到的那一把,验签的一方只认得一把
|
||||
fresh = Ed25519PrivateKey.generate()
|
||||
atomic.create_bytes(path, fresh.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8,
|
||||
serialization.NoEncryption()), mode=0o600)
|
||||
key = serialization.load_pem_private_key(path.read_bytes(), password=None)
|
||||
assert isinstance(key, Ed25519PrivateKey)
|
||||
return key
|
||||
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -62,10 +63,9 @@ async def put_json(
|
||||
artifact_id = content_hash(text)
|
||||
path = _path_of(artifact_id)
|
||||
if not path.exists():
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(".tmp")
|
||||
tmp.write_text(text, encoding="utf-8")
|
||||
tmp.rename(path) # 原子落盘,写一半崩了不会留下半个工件
|
||||
# 原子落盘(core/atomic):写一半崩了不会留下半个工件;临时文件名带随机后缀,两个进程同时写同一份内容
|
||||
# 各写各的,不会一个把另一个写到一半的截断(以前共用 <id>.tmp,后到的会把先到的半成品改名成正式文件)
|
||||
atomic.write_text(path, text)
|
||||
|
||||
async with SessionLocal() as session:
|
||||
existing = await session.get(Artifact, (artifact_id, run_id or ""))
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
"""数据目录里的文件一律先写临时文件(同目录、名字带随机后缀)再改名。
|
||||
|
||||
为什么:读的一方(同一台机器上的别的进程、共享存储另一头的机器)要么看到旧的完整文件,要么看到新的完整文件,
|
||||
看不到写了一半的;进程写到一半没了,留下的是一个没人认的临时文件,不是半个正式文件。临时文件名带随机后缀:
|
||||
两个进程同时写同一个文件(内容寻址的工件撞了同一份内容)各写各的,不会一个把另一个写到一半的截断。
|
||||
|
||||
- write_bytes / write_text:写完 os.replace 到目标(目标已有就替换);
|
||||
- create_bytes / create_text:只在目标还没有时建(os.link,已有就不动、返回 False)——给只该生成一次的东西用
|
||||
(加密密钥、签名私钥):两个进程同时第一次启动,后到的读先到的那一份,不会各拿一把。
|
||||
|
||||
上传表格的版本库、原件存档、记录表另有自己的写法(data/raw_store、data/table_versions,同样先写临时文件),不走这里。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
#: 临时文件名里的记号:.<原名>.tmp-<随机>(同 data/raw_store.TMP_MARK 的写法)
|
||||
TMP_MARK = ".tmp-"
|
||||
|
||||
|
||||
def _tmp_for(path: Path) -> Path:
|
||||
return path.parent / f".{path.name}{TMP_MARK}{secrets.token_hex(6)}"
|
||||
|
||||
|
||||
def _write_tmp(path: Path, data: bytes, *, mode: int | None, fsync: bool) -> Path:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = _tmp_for(path)
|
||||
# 没给 mode 时照普通新建文件(0666 再受 umask 约束);给了的(密钥 0600)从建出来那一刻起就是那个权限,再 chmod
|
||||
# 一次压过 umask
|
||||
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666 if mode is None else 0o600)
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as fh:
|
||||
fh.write(data)
|
||||
fh.flush()
|
||||
if fsync:
|
||||
os.fsync(fh.fileno())
|
||||
if mode is not None:
|
||||
os.chmod(tmp, mode)
|
||||
except BaseException:
|
||||
with contextlib.suppress(OSError):
|
||||
tmp.unlink()
|
||||
raise
|
||||
return tmp
|
||||
|
||||
|
||||
def write_bytes(path: Path, data: bytes, *, mode: int | None = None, fsync: bool = False) -> None:
|
||||
"""原子地写(替换)path。"""
|
||||
tmp = _write_tmp(path, data, mode=mode, fsync=fsync)
|
||||
try:
|
||||
os.replace(tmp, path)
|
||||
except BaseException:
|
||||
with contextlib.suppress(OSError):
|
||||
tmp.unlink()
|
||||
raise
|
||||
|
||||
|
||||
def write_text(path: Path, text: str, *, encoding: str = "utf-8", mode: int | None = None,
|
||||
fsync: bool = False) -> None:
|
||||
write_bytes(path, text.encode(encoding), mode=mode, fsync=fsync)
|
||||
|
||||
|
||||
def create_bytes(path: Path, data: bytes, *, mode: int | None = None, fsync: bool = True) -> bool:
|
||||
"""path 还没有时原子地建出来,返回 True;已经有了(别人先建了)什么都不动,返回 False。"""
|
||||
tmp = _write_tmp(path, data, mode=mode, fsync=fsync)
|
||||
try:
|
||||
os.link(tmp, path)
|
||||
return True
|
||||
except FileExistsError:
|
||||
return False
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
tmp.unlink()
|
||||
|
||||
|
||||
def create_text(path: Path, text: str, *, encoding: str = "utf-8", mode: int | None = None) -> bool:
|
||||
return create_bytes(path, text.encode(encoding), mode=mode)
|
||||
@@ -6,6 +6,7 @@ import os
|
||||
|
||||
from cryptography.fernet import Fernet, InvalidToken
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings
|
||||
|
||||
_PREFIX = "enc::"
|
||||
@@ -16,12 +17,11 @@ def _load_key() -> bytes:
|
||||
raw = settings.secret_key
|
||||
if not raw:
|
||||
key_file = settings.data_dir / ".secret_key"
|
||||
if key_file.exists():
|
||||
raw = key_file.read_text().strip()
|
||||
else:
|
||||
raw = base64.urlsafe_b64encode(os.urandom(32)).decode()
|
||||
key_file.write_text(raw)
|
||||
key_file.chmod(0o600)
|
||||
if not key_file.exists():
|
||||
# 只建一次(core/atomic.create_text):两个进程同时第一次启动,后到的读先到的那一把,不会各拿一把、
|
||||
# 各自加密的东西对方解不开
|
||||
atomic.create_text(key_file, base64.urlsafe_b64encode(os.urandom(32)).decode(), mode=0o600)
|
||||
raw = key_file.read_text().strip()
|
||||
# 任意字符串都归一化成 32 字节 Fernet key
|
||||
digest = hashlib.sha256(raw.encode()).digest()
|
||||
return base64.urlsafe_b64encode(digest)
|
||||
|
||||
@@ -20,6 +20,7 @@ import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.core import atomic
|
||||
from app.files.store import as_ref, bound_files, files_root, read_blob
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -73,10 +74,8 @@ def extract_text(sha: str, name: str, mime: str) -> tuple[str | None, str | None
|
||||
if text is not None and not text.strip():
|
||||
text, problem = None, "没有抽出文字(可能是扫描件)"
|
||||
try:
|
||||
cache.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = cache.with_suffix(".tmp")
|
||||
tmp.write_text(json.dumps({"text": text, "problem": problem}, ensure_ascii=False), encoding="utf-8")
|
||||
tmp.replace(cache)
|
||||
# 临时文件名带随机后缀(core/atomic):两个进程同时抽同一个附件,不会一个把另一个写到一半的改名过去
|
||||
atomic.write_text(cache, json.dumps({"text": text, "problem": problem}, ensure_ascii=False))
|
||||
except OSError:
|
||||
logger.warning("附件 %s 抽出的文字没能缓存", name, exc_info=True)
|
||||
return text, problem
|
||||
|
||||
@@ -11,6 +11,7 @@ import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings, sanitize_session
|
||||
from app.core.errors import describe_exception, raw
|
||||
from app.sandbox import interpreter
|
||||
@@ -176,13 +177,13 @@ class BubblewrapSandbox(Sandbox):
|
||||
entry = entry_name(ext)
|
||||
cmd = [*cmd_prefix, f"/work/{entry}"]
|
||||
workdir = self._session_dir(session_id or f"tmp-{int(time.time() * 1000)}")
|
||||
(workdir / entry).write_text(code)
|
||||
atomic.write_text(workdir / entry, code)
|
||||
for rel, content in (files or {}).items():
|
||||
target = (workdir / rel).resolve()
|
||||
if not str(target).startswith(str(workdir.resolve())):
|
||||
return ExecResult(ok=False, backend=self.name, error=f"非法文件路径:{rel}")
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(content)
|
||||
atomic.write_text(target, content)
|
||||
|
||||
argv = self._argv(workdir, limits, cmd)
|
||||
for key, value in (env or {}).items():
|
||||
|
||||
@@ -9,6 +9,7 @@ import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings, sanitize_session
|
||||
from app.core.errors import describe_exception, raw
|
||||
from app.sandbox import interpreter
|
||||
@@ -107,13 +108,13 @@ class LocalSandbox(Sandbox):
|
||||
)
|
||||
|
||||
workdir = self._session_dir(session_id or f"tmp-{int(time.time()*1000)}")
|
||||
(workdir / entry).write_text(code)
|
||||
atomic.write_text(workdir / entry, code)
|
||||
for rel, content in (files or {}).items():
|
||||
target = (workdir / rel).resolve()
|
||||
if not str(target).startswith(str(workdir.resolve())):
|
||||
return ExecResult(ok=False, backend=self.name, error=f"非法文件路径:{rel}")
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(content)
|
||||
atomic.write_text(target, content)
|
||||
|
||||
# 只透传最小环境变量,避免把宿主机的 API key 泄漏给被执行的代码
|
||||
child_env = {
|
||||
|
||||
@@ -9,6 +9,7 @@ import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import sanitize_session, settings
|
||||
from app.core.errors import describe_exception, raw
|
||||
from app.sandbox import interpreter
|
||||
@@ -194,18 +195,18 @@ class SeatbeltSandbox(Sandbox):
|
||||
ext, cmd_of = _LANG[language]
|
||||
entry = entry_name(ext)
|
||||
workdir = self._session_dir(session_id or f"tmp-{int(time.time() * 1000)}")
|
||||
(workdir / entry).write_text(code)
|
||||
atomic.write_text(workdir / entry, code)
|
||||
|
||||
for rel, content in (files or {}).items():
|
||||
target = (workdir / rel).resolve()
|
||||
if not str(target).startswith(str(workdir.resolve())):
|
||||
return ExecResult(ok=False, backend=self.name, error=f"非法文件路径:{rel}")
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(content)
|
||||
atomic.write_text(target, content)
|
||||
|
||||
self._policies.mkdir(parents=True, exist_ok=True)
|
||||
policy_file = self._policies / f"{workdir.name}-{language}.sb"
|
||||
policy_file.write_text(build_policy(workdir, language=language, network=limits.network))
|
||||
atomic.write_text(policy_file, build_policy(workdir, language=language, network=limits.network))
|
||||
|
||||
argv = [SANDBOX_EXEC, "-f", str(policy_file), *cmd_of(), entry]
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.core import atomic
|
||||
from app.core.config import settings
|
||||
from app.tools.registry import ToolContext, register
|
||||
|
||||
@@ -74,9 +75,10 @@ async def file_write(ctx: ToolContext, **kwargs: Any) -> dict[str, Any]:
|
||||
target = _resolve(args.path, ctx)
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
with target.open("a" if args.append else "w", encoding="utf-8") as f:
|
||||
f.write(args.content)
|
||||
# 先写临时文件再改名(core/atomic):同一个工作目录里正在跑的代码、另一台机器上读共享数据目录的一方,看到的要么是
|
||||
# 旧的整份、要么是新的整份。追加也一样:读出原有的、接上新的、整份换掉
|
||||
before = target.read_bytes() if args.append and target.is_file() else b""
|
||||
atomic.write_bytes(target, before + args.content.encode("utf-8"))
|
||||
return {"path": args.path, "bytes": target.stat().st_size, "ok": True}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
"""数据目录里的文件先写临时文件再改名(core/atomic):读的一方看不到写了一半的,只该生成一次的不会各生成一份。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import stat
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core import atomic
|
||||
from app.db import models # noqa: F401 建表要先认得模型(单独跑这个文件时)
|
||||
|
||||
|
||||
def _leftovers(folder) -> list[str]:
|
||||
return [p.name for p in folder.iterdir() if atomic.TMP_MARK in p.name]
|
||||
|
||||
|
||||
def test_write_replaces_whole_and_leaves_no_temp_file(tmp_path):
|
||||
target = tmp_path / "a" / "b.json"
|
||||
atomic.write_text(target, "第一版")
|
||||
atomic.write_text(target, "第二版")
|
||||
assert target.read_text(encoding="utf-8") == "第二版"
|
||||
assert _leftovers(target.parent) == []
|
||||
|
||||
|
||||
def test_a_failed_write_keeps_the_old_file(tmp_path, monkeypatch):
|
||||
target = tmp_path / "x.txt"
|
||||
atomic.write_text(target, "旧的整份")
|
||||
|
||||
def boom(src, dst):
|
||||
raise OSError("磁盘满了")
|
||||
|
||||
monkeypatch.setattr(os, "replace", boom)
|
||||
with pytest.raises(OSError):
|
||||
atomic.write_text(target, "新的")
|
||||
assert target.read_text(encoding="utf-8") == "旧的整份"
|
||||
assert _leftovers(tmp_path) == []
|
||||
|
||||
|
||||
def test_create_only_once_and_the_loser_reads_the_winner(tmp_path):
|
||||
"""几个线程同时第一次建同一个密钥文件:只有一个建成,文件内容就是那一个写的,没有临时文件留下。"""
|
||||
target = tmp_path / ".secret"
|
||||
results: list[tuple[str, bool]] = []
|
||||
barrier = threading.Barrier(8)
|
||||
|
||||
def make(i: int) -> None:
|
||||
barrier.wait()
|
||||
results.append((f"key-{i}", atomic.create_text(target, f"key-{i}", mode=0o600)))
|
||||
|
||||
threads = [threading.Thread(target=make, args=(i,)) for i in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
winners = [k for k, made in results if made]
|
||||
assert len(winners) == 1 and target.read_text() == winners[0]
|
||||
assert stat.S_IMODE(target.stat().st_mode) == 0o600
|
||||
assert _leftovers(tmp_path) == []
|
||||
|
||||
|
||||
def test_content_addressed_artifacts_are_written_atomically(tmp_path, monkeypatch):
|
||||
"""工件按内容寻址:同一份内容写两遍,文件只有一份、内容完整,不留临时文件。"""
|
||||
import asyncio
|
||||
|
||||
from app.core import artifact_store
|
||||
from app.core.config import settings
|
||||
|
||||
async def twice() -> tuple[str, str]:
|
||||
a = await artifact_store.put_json({"区域": "虚构", "值": 1}, kind="test_atomic")
|
||||
b = await artifact_store.put_json({"值": 1, "区域": "虚构"}, kind="test_atomic")
|
||||
return a, b
|
||||
|
||||
a, b = asyncio.run(twice())
|
||||
assert a == b and artifact_store.load(a) == {"区域": "虚构", "值": 1}
|
||||
folder = settings.data_dir / "artifacts" / a[:2]
|
||||
assert _leftovers(folder) == []
|
||||
|
||||
|
||||
async def test_file_write_tool_appends_by_replacing_the_whole_file():
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.tools.builtin.files import file_write
|
||||
|
||||
ctx = SimpleNamespace(sandbox_session=f"atomic-{os.getpid()}")
|
||||
assert (await file_write(ctx, path="notes/a.txt", content="第一行\n"))["ok"]
|
||||
out = await file_write(ctx, path="notes/a.txt", content="第二行\n", append=True)
|
||||
from app.core.config import settings
|
||||
|
||||
path = settings.session_dir(ctx.sandbox_session) / "notes" / "a.txt"
|
||||
assert path.read_text(encoding="utf-8") == "第一行\n第二行\n" and out["bytes"] == len("第一行\n第二行\n".encode())
|
||||
assert _leftovers(path.parent) == []
|
||||
Reference in New Issue
Block a user