diff --git a/.agents/MEMORY.md b/.agents/MEMORY.md index f732316..497ece4 100644 --- a/.agents/MEMORY.md +++ b/.agents/MEMORY.md @@ -5,6 +5,7 @@ ## 用户偏好与长期原则 - 一律使用中文与用户沟通。 +- Git commit message 一律使用中文。 - Windows Python 脚本 stdout 可能是 GBK:使用 ASCII 状态标签,不用 emoji 或特殊装饰字符。 - `CHANGELOG.md` 提到海外模型时只用“国际旗舰模型”等泛称;具体型号只放在 `PROGRESS.md`、配置和 git log。 - 版本、CHANGELOG、PROGRESS 在 push 前统一更新;DESIGN 跟随产生架构或决策变化的 commit。 @@ -61,4 +62,3 @@ Mermaid/Chromium 的历史故障最终有三个根因: ## 领域 用户单位是中国建筑材料科学研究总院。代码、库、模板和示例默认服务于水泥/混凝土、玻璃、陶瓷、耐火和新型建材的材料研发、表征分析、实验建模与科研写作;不是建筑施工、BIM 或结构设计语境。 - diff --git a/DESIGN.md b/DESIGN.md index 10ffe66..40339cd 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -29,6 +29,8 @@ zcbot/ │ ├── skills.py # SkillRegistry(渐进披露,多来源) │ ├── task.py # TaskState │ ├── memory.py # per-user .memory/ 双层记忆 +│ ├── file_store.py # 原子文件替换 + 跨进程 advisory file lock +│ ├── kb_lock.py # .kb/<库> mutation 单写者锁(Web/入库/fs tool 共用) │ ├── shortcuts.py # 快捷指令(入口层确定性展开) │ ├── paths.py # task_dir db form 归一 │ ├── storage/ # SQLAlchemy 2.x ORM;usage(计费写)/telemetry(失败埋点)/usage_report(聚合读)三分 @@ -94,7 +96,7 @@ Session = 消息列表,ORM 直写 PG `messages`(append-only,jsonb 存 LiteLLM 用户自建资料(规范/报告/标准/内部文档)的长期查阅层,与 §3.7 记忆同范式:**纯文件 + prompt 注入契约,无向量无 DB**(判据同"真实文件为准":个人库几十~百余文件,agentic search 足够;索引若引入只能是可重建派生缓存)。**两层格局**:zcbot 内建 `.kb/` 私有小库(本节)+ 院检索服务共享大库(document_search,zcbot 只当客户端)——分工标准 = 文件数 × 查询频次,路由靠各自工具/契约描述自然分流,不在 kb 契约里点名 document_search(2026-07-22 收窄:契约只管自己怎么用,少一层耦合)。 - **做成机制而非 skill**(判据:有独立于会话的持久状态需用户管理 → 机制):落盘 `user_root/.kb/<库名>/`(INDEX.md + docs/ 转换后 md + sources/ 原件)。**"已入库"判据 = INDEX.md 有条目**,sources 有而 INDEX 无 = 待入库 → 入库幂等、崩溃可恢复、零 migration。dotfile 命名同 `.memory` 双向防呆,GET /v1/files 天然隐藏。 -- **入库管线**(`core/kb_ingest.py`,上传即触发 + 手动兜底):markitdown Python API 转 md → 扫描件 PDF(文本近零)走方舟文档理解 OCR 兜底(§8.13 同通道)→ deepseek flash 单次 chat 写 标题/摘要/关键词(失败降级文件名+正文开头,不阻塞)→ 追加 INDEX 行。编排照定时执行器:create_task + to_thread + per-(user,库) 内存锁去重;进度存内存供前端轮询,崩了靠 FS 判据续跑。 +- **入库管线**(`core/kb_ingest.py`,上传即触发 + 手动兜底):markitdown Python API 转 md → 扫描件 PDF(文本近零)走方舟文档理解 OCR 兜底(§8.13 同通道)→ deepseek flash 单次 chat 写 标题/摘要/关键词(失败降级文件名+正文开头,不阻塞)→ 追加 INDEX 行。编排照定时执行器:create_task + to_thread;**写并发收口为共享 FS advisory lock**(`.kb/.locks/.lock`):Web 上传/删除、后台入库与 agent `write/edit` 对同一库共用一把跨进程锁,蓝绿实例间只允许一个写者,锁占用返 409/工具可重试;进程退出由 OS 自动释锁,不靠清理锁文件。文档正文、原件与 INDEX 全走同目录临时文件 + fsync + `os.replace` 原子发布,读者只会看到完整旧版或完整新版。进度详情仍以内存保存细节,但会探测跨进程锁补出 `running`;崩了靠 FS 判据续跑。 - **agent 侧零新工具**:`kb_block`(照 memory_block)把 INDEX 全文 + 契约(主动查阅无需点名 / INDEX 行格式 / 答题标来源)注 prompt;**零库时注极简冷启动契约**(建库步骤 + INDEX 行格式 + "成篇资料进 KB / 短事实进记忆"分工,~百 token)——原"有库才注入"省 token,但零注入让模型不知道 KB 机制存在,用户说"放进知识库"被就近写进 `.memory/`(2026-07-24 真实事故),与 memory 空契约常驻是同一课:教会第一次,建库落盘后下轮 build_agent 自然切全量。fs 工具在 user_root 内可读写、docker 沙箱整 user_root bind → `.kb` 天然可达。INDEX 行格式是对话内手动入库与后台产出的同一契约。 - **API 薄壳**(`/v1/kb*` 8 端点):列/建/删库、详情(带入库进度)、上传即入库、手动 ingest、看/删单篇。**不设 HTTP 检索端点**——检索是 agent 的事。前端两栏 modal(kb.js)管上传/删除,查询全走对话。 - **记账**:`usage_events` kind="kb_ingest"(OCR 那笔走 kind="vision"),无 task 上下文 → 0022 放宽 task_id 可 NULL,溯源靠 units JSONB `{"kb", "source"}`。 diff --git a/core/file_store.py b/core/file_store.py new file mode 100644 index 0000000..fe7b9ed --- /dev/null +++ b/core/file_store.py @@ -0,0 +1,135 @@ +"""Durable local-file primitives shared by host and sandbox file mutations. + +Writes are staged in the destination directory, fsynced, then published with +``os.replace`` so readers see either the old complete file or the new complete +file. ``interprocess_file_lock`` uses the operating system's advisory lock; +the lock file may remain on disk, but the lock itself is released automatically +when a process exits. +""" +from __future__ import annotations + +import os +import tempfile +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator, Optional + + +class FileLockBusy(RuntimeError): + """A non-blocking or timed inter-process lock could not be acquired.""" + + +def _atomic_replace(path: Path, data: bytes) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + old_mode: Optional[int] = None + try: + old_mode = path.stat().st_mode + except FileNotFoundError: + pass + + fd, raw_tmp = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent) + ) + tmp = Path(raw_tmp) + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + if old_mode is not None: + os.chmod(tmp, old_mode) + os.replace(tmp, path) + # Persist the directory entry on POSIX. Windows cannot open directories + # this way; os.replace still gives atomic visibility there. + if os.name != "nt": + try: + dir_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(dir_fd) + finally: + os.close(dir_fd) + except OSError: + # Some network/virtual filesystems reject directory fsync. + # The file itself is already fsynced and atomically visible. + pass + finally: + try: + tmp.unlink() + except FileNotFoundError: + pass + + +def atomic_write_bytes(path: Path, data: bytes) -> None: + _atomic_replace(Path(path), data) + + +def atomic_write_text(path: Path, text: str, encoding: str = "utf-8") -> None: + _atomic_replace(Path(path), text.encode(encoding)) + + +def _try_lock(f) -> bool: + f.seek(0) + if os.name == "nt": + import msvcrt + + try: + msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK, 1) + return True + except OSError: + return False + + import fcntl + + try: + fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + return True + except BlockingIOError: + return False + + +def _unlock(f) -> None: + f.seek(0) + if os.name == "nt": + import msvcrt + + msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1) + return + + import fcntl + + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + + +@contextmanager +def interprocess_file_lock( + path: Path, + *, + timeout_seconds: Optional[float] = 0, + poll_seconds: float = 0.05, +) -> Iterator[None]: + """Acquire an advisory exclusive lock. + + ``timeout_seconds=0`` is non-blocking; ``None`` waits indefinitely. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a+b") as f: + f.seek(0, os.SEEK_END) + if f.tell() == 0: + f.write(b"\0") + f.flush() + + deadline = ( + None if timeout_seconds is None + else time.monotonic() + max(0.0, timeout_seconds) + ) + while not _try_lock(f): + if deadline is not None and time.monotonic() >= deadline: + raise FileLockBusy(str(path)) + time.sleep(poll_seconds) + try: + yield + finally: + _unlock(f) diff --git a/core/kb.py b/core/kb.py index adf44e2..2871b3f 100644 --- a/core/kb.py +++ b/core/kb.py @@ -21,9 +21,12 @@ from __future__ import annotations import re import shutil from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Iterable, List, Optional from uuid import UUID +from .file_store import atomic_write_bytes, atomic_write_text +from .kb_lock import kb_mutation_lock + # INDEX 单行格式(全角 | 分隔,摘要/关键词内允许半角标点)。agent 对话内手动入库 # 与后台 ingest 产出同一格式 —— 契约文本(kb_block)里原样给出。 INDEX_LINE_FORMAT = "- [标题](docs/<文件名>.md)|来源 sources/<原件名>|摘要:<两三句>|关键词:<逗号分隔>" @@ -143,11 +146,12 @@ def create_kb(workspace_dir: Path, user_id: UUID, name: str) -> Optional[Path]: d = kb_dir(workspace_dir, user_id, name) if d is None: return None - (d / "docs").mkdir(parents=True, exist_ok=True) - (d / "sources").mkdir(parents=True, exist_ok=True) - idx = d / "INDEX.md" - if not idx.exists(): - idx.write_text(f"# {name}\n\n", encoding="utf-8") + with kb_mutation_lock(workspace_dir, user_id, name): + (d / "docs").mkdir(parents=True, exist_ok=True) + (d / "sources").mkdir(parents=True, exist_ok=True) + idx = d / "INDEX.md" + if not idx.exists(): + atomic_write_text(idx, f"# {name}\n\n") return d @@ -156,7 +160,10 @@ def delete_kb(workspace_dir: Path, user_id: UUID, name: str) -> bool: d = kb_dir(workspace_dir, user_id, name) if d is None or not d.is_dir(): return False - shutil.rmtree(d) + with kb_mutation_lock(workspace_dir, user_id, name): + if not d.is_dir(): + return False + shutil.rmtree(d) return True @@ -186,11 +193,10 @@ def read_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> Op return None -def delete_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> bool: - """删单篇:docs 文件 + INDEX 对应行 + 对应 source 原件一起删(否则原件会被当 - 待入库重新转一遍 —— 判据使然)。""" - d = kb_dir(workspace_dir, user_id, name) - if d is None or not is_safe_file_name(filename) or not filename.endswith(".md"): +def _delete_doc_unlocked( + d: Path, filename: str, *, delete_source: bool = True +) -> bool: + if not is_safe_file_name(filename) or not filename.endswith(".md"): return False doc_rel = f"docs/{filename}" entries = _read_index(d) @@ -200,14 +206,11 @@ def delete_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> return False if not target.is_file() and hit is None: return False - if target.is_file(): - target.unlink() + src: Optional[Path] = None if hit is not None: src_name = hit["source"].removeprefix("sources/") if is_safe_file_name(src_name): src = d / "sources" / src_name - if src.is_file(): - src.unlink() idx = d / "INDEX.md" try: lines = idx.read_text(encoding="utf-8").splitlines() @@ -215,28 +218,74 @@ def delete_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> ln for ln in lines if not (_INDEX_LINE_RE.match(ln.strip()) and f"({doc_rel})" in ln) ] - idx.write_text("\n".join(kept).rstrip() + "\n", encoding="utf-8") + atomic_write_text(idx, "\n".join(kept).rstrip() + "\n") except (OSError, UnicodeDecodeError): - pass + return False + # Publish the recoverable INDEX state before deleting data:if the process + # dies here, an undeleted source is merely pending and can be re-ingested. + if target.is_file(): + target.unlink() + if delete_source and src is not None and src.is_file(): + src.unlink() return True -def save_source(workspace_dir: Path, user_id: UUID, name: str, filename: str, data: bytes) -> Optional[str]: - """上传原件落 sources/(同名覆盖 —— 重传即重新入库的自然语义)。 - 返回落盘文件名;库不存在 / 文件名非法 → None。""" +def delete_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> bool: + """删单篇:docs 文件 + INDEX 对应行 + 对应 source 原件一起删。""" d = kb_dir(workspace_dir, user_id, name) - if d is None or not d.is_dir() or not is_safe_file_name(filename): + if d is None: + return False + with kb_mutation_lock(workspace_dir, user_id, name): + return _delete_doc_unlocked(d, filename) + + +def _save_source_unlocked(d: Path, filename: str, data: bytes) -> Optional[str]: + if not d.is_dir() or not is_safe_file_name(filename): return None src_dir = d / "sources" src_dir.mkdir(parents=True, exist_ok=True) # 覆盖旧 doc 判据:同名 source 若已在 INDEX,删掉旧条目让它重新排队入库 doc_rel_hits = [e for e in _read_index(d) if e["source"] == f"sources/{filename}"] for e in doc_rel_hits: - delete_doc(workspace_dir, user_id, name, e["doc"].removeprefix("docs/")) - (src_dir / filename).write_bytes(data) + if not _delete_doc_unlocked( + d, e["doc"].removeprefix("docs/"), delete_source=False + ): + return None + atomic_write_bytes(src_dir / filename, data) return filename +def save_sources( + workspace_dir: Path, + user_id: UUID, + name: str, + items: Iterable[tuple[str, bytes]], +) -> Optional[List[str]]: + """一把库锁内批量保存原件;任一文件名非法时整批不写。""" + d = kb_dir(workspace_dir, user_id, name) + batch = list(items) + if ( + d is None + or not d.is_dir() + or any(not is_safe_file_name(filename) for filename, _data in batch) + ): + return None + with kb_mutation_lock(workspace_dir, user_id, name): + saved: List[str] = [] + for filename, data in batch: + result = _save_source_unlocked(d, filename, data) + if result is None: + return None + saved.append(result) + return saved + + +def save_source(workspace_dir: Path, user_id: UUID, name: str, filename: str, data: bytes) -> Optional[str]: + """上传单个原件;保留给内部调用方的兼容入口。""" + saved = save_sources(workspace_dir, user_id, name, [(filename, data)]) + return saved[0] if saved else None + + # ── prompt 注入(照 memory_block 范式) ──────────────────────────────── # 零库时的冷启动契约:曾因零库零注入,用户说「放进知识库」被 agent 就近写进 diff --git a/core/kb_ingest.py b/core/kb_ingest.py index 0731a2d..8d8de57 100644 --- a/core/kb_ingest.py +++ b/core/kb_ingest.py @@ -19,7 +19,6 @@ from __future__ import annotations import base64 import hashlib import re -import threading from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -28,6 +27,8 @@ from uuid import UUID from core.ark_client import ArkClient, ArkError from core.capabilities import ModelCapabilities from core.kb import format_index_line, kb_dir, parse_index, pending_sources +from core.file_store import atomic_write_text +from core.kb_lock import KbBusyError, kb_is_locked, kb_mutation_lock from core.llm import LLM from core.storage.usage import record_chat_usage, record_vision_usage @@ -54,20 +55,18 @@ _SUMMARY_PROMPT = """\ {content}""" # ── 进度状态(内存,供 API 轮询) ───────────────────────────────────── -_guard = threading.Lock() -_locks: Dict[Tuple[str, str], threading.Lock] = {} _status: Dict[Tuple[str, str], Dict[str, Any]] = {} -def _lock_for(key: Tuple[str, str]) -> threading.Lock: - with _guard: - return _locks.setdefault(key, threading.Lock()) - - -def ingest_status(user_id: UUID, name: str) -> Dict[str, Any]: +def ingest_status( + user_id: UUID, name: str, workspace_dir: Optional[Path] = None +) -> Dict[str, Any]: """当前/最近一次入库进度(无记录返回 idle);前端轮询 + 库详情附带。""" st = _status.get((str(user_id), name)) - return dict(st) if st else {"running": False, "total": 0, "done": 0, "errors": []} + out = dict(st) if st else {"running": False, "total": 0, "done": 0, "errors": []} + if workspace_dir is not None and not out.get("running"): + out["running"] = kb_is_locked(workspace_dir, user_id, name) + return out # ── 转换 ───────────────────────────────────────────────────────────── @@ -227,7 +226,7 @@ def _append_index(d: Path, line: str) -> None: text = idx.read_text(encoding="utf-8") if idx.is_file() else "" if text and not text.endswith("\n"): text += "\n" - idx.write_text(text + line + "\n", encoding="utf-8") + atomic_write_text(idx, text + line + "\n") def run_ingest(workspace_dir: Path, user_id: UUID, kb_name: str, models_dir: Path) -> bool: @@ -240,9 +239,20 @@ def run_ingest(workspace_dir: Path, user_id: UUID, kb_name: str, models_dir: Pat if d is None or not d.is_dir(): return False key = (str(user_id), kb_name) - lock = _lock_for(key) - if not lock.acquire(blocking=False): + try: + with kb_mutation_lock(workspace_dir, user_id, kb_name): + return _run_ingest_locked(d, key, user_id, kb_name, models_dir) + except KbBusyError: return False + + +def _run_ingest_locked( + d: Path, + key: Tuple[str, str], + user_id: UUID, + kb_name: str, + models_dir: Path, +) -> bool: try: pending = pending_sources(d) st: Dict[str, Any] = { @@ -273,7 +283,7 @@ def run_ingest(workspace_dir: Path, user_id: UUID, kb_name: str, models_dir: Pat raise ValueError(f"抽出文本过少({len(text)} 字符),无法入库") doc_name = _doc_name_for(d, name) (d / "docs").mkdir(parents=True, exist_ok=True) - (d / "docs" / doc_name).write_text(text, encoding="utf-8") + atomic_write_text(d / "docs" / doc_name, text) if llm is not None and caps is not None: title, summary, keywords = _summarize( llm, caps, text=text, filename=name, user_id=user_id, kb_name=kb_name @@ -296,4 +306,3 @@ def run_ingest(workspace_dir: Path, user_id: UUID, kb_name: str, models_dir: Pat st["running"] = False st["current"] = None st["finished_at"] = datetime.now().isoformat(timespec="seconds") - lock.release() diff --git a/core/kb_lock.py b/core/kb_lock.py new file mode 100644 index 0000000..7c56ef8 --- /dev/null +++ b/core/kb_lock.py @@ -0,0 +1,81 @@ +"""Cross-process mutation lock for one user's knowledge-base library.""" +from __future__ import annotations + +import hashlib +from contextlib import contextmanager, nullcontext +from pathlib import Path +from typing import Iterator, Optional +from uuid import UUID + +from .file_store import FileLockBusy, interprocess_file_lock + + +class KbBusyError(RuntimeError): + """The target knowledge base is being mutated by another process.""" + + +def _user_root(workspace_dir: Path, user_id: UUID) -> Path: + return Path(workspace_dir) / "users" / str(user_id) + + +def kb_lock_path(user_root: Path, kb_name: str) -> Path: + digest = hashlib.sha256(kb_name.encode("utf-8")).hexdigest()[:32] + return Path(user_root) / ".kb" / ".locks" / f"{digest}.lock" + + +@contextmanager +def kb_mutation_lock( + workspace_dir: Path, + user_id: UUID, + kb_name: str, + *, + timeout_seconds: Optional[float] = 0, +) -> Iterator[None]: + try: + with interprocess_file_lock( + kb_lock_path(_user_root(workspace_dir, user_id), kb_name), + timeout_seconds=timeout_seconds, + ): + yield + except FileLockBusy as e: + raise KbBusyError(kb_name) from e + + +@contextmanager +def kb_mutation_lock_for_path( + target: Path, + user_root: Optional[Path], + *, + timeout_seconds: Optional[float] = 0, +) -> Iterator[None]: + """Lock the containing ``.kb/`` library; no-op outside ``.kb``.""" + if user_root is None: + with nullcontext(): + yield + return + try: + rel = Path(target).resolve().relative_to(Path(user_root).resolve()) + except (OSError, ValueError): + with nullcontext(): + yield + return + if len(rel.parts) < 3 or rel.parts[0] != ".kb" or rel.parts[1].startswith("."): + with nullcontext(): + yield + return + try: + with interprocess_file_lock( + kb_lock_path(Path(user_root), rel.parts[1]), + timeout_seconds=timeout_seconds, + ): + yield + except FileLockBusy as e: + raise KbBusyError(rel.parts[1]) from e + + +def kb_is_locked(workspace_dir: Path, user_id: UUID, kb_name: str) -> bool: + try: + with kb_mutation_lock(workspace_dir, user_id, kb_name): + return False + except KbBusyError: + return True diff --git a/core/session.py b/core/session.py index 456988a..8b0265b 100644 --- a/core/session.py +++ b/core/session.py @@ -10,7 +10,6 @@ system prompt **不入库** —— 每次 build_agent 重建拼到 messages[0]( """ from __future__ import annotations -import os from pathlib import Path from typing import Any, Dict, List, Optional from uuid import UUID @@ -19,6 +18,7 @@ from sqlalchemy import delete, func, select from .storage import session_scope from .storage.models import Message, Task +from .file_store import atomic_write_text def _to_dict(msg: Any) -> Any: @@ -31,21 +31,6 @@ def _to_dict(msg: Any) -> Any: return msg -def atomic_write_text(path: Path, text: str, encoding: str = "utf-8") -> None: - """原子写: 先写到 path.tmp 再 os.replace 到 path。 - - 防止写中途异常(磁盘满 / surrogate 编码错 / 进程被杀)留下 0 字节或半文件。 - skill 产物(*.spec.md / sections/*.md 等)走这里,messages 已改走 PG。 - """ - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_suffix(path.suffix + ".tmp") - with open(tmp, "w", encoding=encoding, newline="\n") as f: - f.write(text) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - - class Session: """消息列表 anchored on task_id。 diff --git a/deploy/sandbox/Dockerfile b/deploy/sandbox/Dockerfile index 245a868..f2e6df7 100644 --- a/deploy/sandbox/Dockerfile +++ b/deploy/sandbox/Dockerfile @@ -172,6 +172,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends file \ # edit/glob/grep 全在容器内执行,物理边界替代代码护栏。tools/ 目录与 host 同步 # (build 时 COPY,不挂 mount ── 容器内代码不应跟随 host repo 修改重启)。 COPY tools/ /sandbox/tools/ +COPY core/__init__.py core/file_store.py core/kb_lock.py /sandbox/core/ COPY core/sandbox/tool_runner.py /sandbox/tool_runner.py COPY deploy/sandbox/init.sh /init.sh diff --git a/tests/test_file_store.py b/tests/test_file_store.py new file mode 100644 index 0000000..de7f025 --- /dev/null +++ b/tests/test_file_store.py @@ -0,0 +1,175 @@ +"""原子文件写与知识库跨进程写锁回归。""" +from __future__ import annotations + +import tempfile +import unittest +from multiprocessing import get_context +from pathlib import Path +from unittest.mock import patch +from uuid import UUID + +from core.file_store import ( + FileLockBusy, + atomic_write_text, + interprocess_file_lock, +) +from core.kb import ( + create_kb, + delete_doc, + delete_kb, + format_index_line, + save_source, + save_sources, +) +from core.kb_lock import KbBusyError, kb_is_locked, kb_mutation_lock +from tools.fs import WriteTool + + +_UID = UUID("11111111-2222-3333-4444-555555555555") + + +def _hold_lock_in_child(path: str, ready, release) -> None: + with interprocess_file_lock(Path(path)): + ready.set() + release.wait(10) + + +class AtomicWriteTests(unittest.TestCase): + def test_atomic_write_replaces_complete_file_and_cleans_temp(self): + with tempfile.TemporaryDirectory() as td: + path = Path(td) / "a.txt" + path.write_text("old", encoding="utf-8") + atomic_write_text(path, "新内容") + self.assertEqual(path.read_text(encoding="utf-8"), "新内容") + self.assertEqual(list(path.parent.glob(".a.txt.*.tmp")), []) + + def test_replace_failure_keeps_old_file_and_cleans_temp(self): + with tempfile.TemporaryDirectory() as td: + path = Path(td) / "a.txt" + path.write_text("old", encoding="utf-8") + with patch("core.file_store.os.replace", side_effect=OSError("boom")): + with self.assertRaises(OSError): + atomic_write_text(path, "new") + self.assertEqual(path.read_text(encoding="utf-8"), "old") + self.assertEqual(list(path.parent.glob(".a.txt.*.tmp")), []) + + +class FileLockTests(unittest.TestCase): + def test_second_handle_cannot_acquire_locked_file(self): + with tempfile.TemporaryDirectory() as td: + lock = Path(td) / "x.lock" + with interprocess_file_lock(lock): + with self.assertRaises(FileLockBusy): + with interprocess_file_lock(lock): + pass + with interprocess_file_lock(lock): + pass + + def test_lock_is_visible_across_processes(self): + with tempfile.TemporaryDirectory() as td: + lock = Path(td) / "cross-process.lock" + ctx = get_context("spawn") + ready = ctx.Event() + release = ctx.Event() + child = ctx.Process( + target=_hold_lock_in_child, args=(str(lock), ready, release) + ) + child.start() + try: + self.assertTrue(ready.wait(5), "child did not acquire lock") + with self.assertRaises(FileLockBusy): + with interprocess_file_lock(lock): + pass + finally: + release.set() + child.join(5) + if child.is_alive(): + child.terminate() + child.join(5) + self.assertEqual(child.exitcode, 0) + + +class SandboxPackagingTests(unittest.TestCase): + def test_sandbox_copies_file_store_dependencies(self): + root = Path(__file__).resolve().parents[1] + dockerfile = (root / "deploy" / "sandbox" / "Dockerfile").read_text("utf-8") + self.assertIn("core/file_store.py", dockerfile) + self.assertIn("core/kb_lock.py", dockerfile) + + +class KnowledgeBaseMutationTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.ws = Path(self.tmp.name) + self.user_root = self.ws / "users" / str(_UID) + self.kb = create_kb(self.ws, _UID, "标准库") + assert self.kb is not None + + def tearDown(self): + self.tmp.cleanup() + + def test_lock_blocks_service_and_agent_mutations(self): + idx = self.kb / "INDEX.md" + old = idx.read_text(encoding="utf-8") + tool = WriteTool(base_dir=self.user_root, user_root=self.user_root) + + with kb_mutation_lock(self.ws, _UID, "标准库"): + self.assertTrue(kb_is_locked(self.ws, _UID, "标准库")) + with self.assertRaises(KbBusyError): + delete_kb(self.ws, _UID, "标准库") + result = tool.execute(".kb/标准库/INDEX.md", "bad") + self.assertTrue(result.startswith("[Error]")) + self.assertEqual(idx.read_text(encoding="utf-8"), old) + + self.assertFalse(kb_is_locked(self.ws, _UID, "标准库")) + result = tool.execute(".kb/标准库/INDEX.md", "ok") + self.assertTrue(result.startswith("[wrote")) + self.assertEqual(idx.read_text(encoding="utf-8"), "ok") + + def test_invalid_batch_is_rejected_before_any_write(self): + result = save_sources( + self.ws, + _UID, + "标准库", + [("good.txt", b"good"), ("../bad.txt", b"bad")], + ) + self.assertIsNone(result) + self.assertFalse((self.kb / "sources" / "good.txt").exists()) + + def _seed_indexed_doc(self): + source = self.kb / "sources" / "a.txt" + doc = self.kb / "docs" / "a.md" + source.write_bytes(b"old") + doc.write_text("old doc", encoding="utf-8") + line = format_index_line( + title="A", + doc="docs/a.md", + source="sources/a.txt", + summary="old", + keywords="a", + ) + (self.kb / "INDEX.md").write_text( + f"# 标准库\n\n{line}\n", encoding="utf-8" + ) + return source, doc + + def test_overwrite_removes_old_index_and_doc_then_publishes_source(self): + source, doc = self._seed_indexed_doc() + self.assertEqual( + save_source(self.ws, _UID, "标准库", "a.txt", b"new"), "a.txt" + ) + self.assertEqual(source.read_bytes(), b"new") + self.assertFalse(doc.exists()) + self.assertNotIn("docs/a.md", (self.kb / "INDEX.md").read_text("utf-8")) + + def test_delete_keeps_data_when_index_publish_fails(self): + source, doc = self._seed_indexed_doc() + with patch("core.kb.atomic_write_text", side_effect=OSError("disk full")): + self.assertFalse(delete_doc(self.ws, _UID, "标准库", "a.md")) + self.assertTrue(source.exists()) + self.assertTrue(doc.exists()) + self.assertIn("docs/a.md", (self.kb / "INDEX.md").read_text("utf-8")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/fs.py b/tools/fs.py index 5c176bd..2856edc 100644 --- a/tools/fs.py +++ b/tools/fs.py @@ -9,6 +9,8 @@ import re from pathlib import Path from .base import Tool +from core.file_store import atomic_write_text +from core.kb_lock import KbBusyError, kb_mutation_lock_for_path class ReadTool(Tool): @@ -64,8 +66,11 @@ class WriteTool(Tool): def execute(self, path: str, content: str) -> str: p = self._resolve(path) - p.parent.mkdir(parents=True, exist_ok=True) - p.write_text(content, encoding="utf-8") + try: + with kb_mutation_lock_for_path(p, self.user_root): + atomic_write_text(p, content) + except KbBusyError: + return "[Error] knowledge base is being updated; wait and retry" return f"[wrote {len(content)} chars to {self._display(p)}]" @@ -90,13 +95,17 @@ class EditTool(Tool): disp = self._display(p) if not p.exists(): return f"[Error] file not found: {disp}" - content = p.read_text(encoding="utf-8") - count = content.count(old_str) - if count == 0: - return f"[Error] old_str not found in {disp}" - if count > 1: - return f"[Error] old_str appears {count} times in {disp}, must be unique — add more context" - p.write_text(content.replace(old_str, new_str), encoding="utf-8") + try: + with kb_mutation_lock_for_path(p, self.user_root): + content = p.read_text(encoding="utf-8") + count = content.count(old_str) + if count == 0: + return f"[Error] old_str not found in {disp}" + if count > 1: + return f"[Error] old_str appears {count} times in {disp}, must be unique — add more context" + atomic_write_text(p, content.replace(old_str, new_str)) + except KbBusyError: + return "[Error] knowledge base is being updated; wait and retry" return f"[edited {disp}: 1 replacement]" diff --git a/web/routers/kb.py b/web/routers/kb.py index 860929d..2aeed6b 100644 --- a/web/routers/kb.py +++ b/web/routers/kb.py @@ -43,9 +43,13 @@ def register_kb_routes(app, *, require_user) -> None: """建库(幂等)。名字非法 → 400。""" from core.agent_builder import resolve_workspace from core.kb import create_kb + from core.kb_lock import KbBusyError name = (body.name or "").strip() - if create_kb(resolve_workspace(None), user_id, name) is None: - raise HTTPException(400, f"invalid kb name: {body.name!r}") + try: + if create_kb(resolve_workspace(None), user_id, name) is None: + raise HTTPException(400, f"invalid kb name: {body.name!r}") + except KbBusyError: + raise HTTPException(409, "该库正在被其他操作修改") return {"name": name} @app.get("/v1/kb/{name}", tags=["kb"]) @@ -54,10 +58,11 @@ def register_kb_routes(app, *, require_user) -> None: from core.agent_builder import resolve_workspace from core.kb import kb_detail from core.kb_ingest import ingest_status - detail = kb_detail(resolve_workspace(None), user_id, name) + ws = resolve_workspace(None) + detail = kb_detail(ws, user_id, name) if detail is None: raise HTTPException(404, f"kb not found: {name!r}") - detail["ingest"] = ingest_status(user_id, name) + detail["ingest"] = ingest_status(user_id, name, ws) return detail @app.delete("/v1/kb/{name}", tags=["kb"]) @@ -65,11 +70,12 @@ def register_kb_routes(app, *, require_user) -> None: """整库删除(原件 + docs + INDEX)。入库进行中 → 409(避免半截写盘)。""" from core.agent_builder import resolve_workspace from core.kb import delete_kb - from core.kb_ingest import ingest_status - if ingest_status(user_id, name).get("running"): - raise HTTPException(409, "该库正在入库,等入库结束再删除") - if not delete_kb(resolve_workspace(None), user_id, name): - raise HTTPException(404, f"kb not found: {name!r}") + from core.kb_lock import KbBusyError + try: + if not delete_kb(resolve_workspace(None), user_id, name): + raise HTTPException(404, f"kb not found: {name!r}") + except KbBusyError: + raise HTTPException(409, "该库正在入库或被其他操作修改") return {"deleted": name} @app.post("/v1/kb/{name}/upload", tags=["kb"]) @@ -84,7 +90,8 @@ def register_kb_routes(app, *, require_user) -> None: 磁盘配额 gate 与 /v1/files/upload 同款。 """ from core.agent_builder import load_config as _load_cfg, resolve_workspace - from core.kb import kb_dir, save_source + from core.kb import kb_dir, save_sources + from core.kb_lock import KbBusyError from core.storage.disk_quota import check_disk_quota, parse_bytes _quotas_cfg = (_load_cfg().get("quotas") or {}) _limit = parse_bytes(_quotas_cfg.get("disk_bytes_per_user")) @@ -97,16 +104,23 @@ def register_kb_routes(app, *, require_user) -> None: d = kb_dir(ws, user_id, name) if d is None or not d.is_dir(): raise HTTPException(404, f"kb not found: {name!r}") - saved: list[dict] = [] + uploads: list[tuple[str, bytes]] = [] for up in files or []: raw_name = up.filename or "" data = await up.read() - ok = save_source(ws, user_id, name, raw_name, data) - if ok is None: - raise HTTPException(400, f"invalid filename: {raw_name!r}") - saved.append({"name": ok, "size": len(data)}) - if not saved: + uploads.append((raw_name, data)) + if not uploads: raise HTTPException(400, "no files uploaded") + try: + saved_names = save_sources(ws, user_id, name, uploads) + except KbBusyError: + raise HTTPException(409, "该库正在入库或被其他操作修改") + if saved_names is None: + raise HTTPException(400, "存在非法文件名") + saved = [ + {"name": saved_name, "size": len(data)} + for saved_name, (_raw_name, data) in zip(saved_names, uploads) + ] _spawn_kb_ingest(user_id, name) return {"count": len(saved), "saved": saved} @@ -135,6 +149,10 @@ def register_kb_routes(app, *, require_user) -> None: """删单篇(docs 文件 + INDEX 行 + source 原件,防原件被重新入库)。""" from core.agent_builder import resolve_workspace from core.kb import delete_doc - if not delete_doc(resolve_workspace(None), user_id, name, filename): - raise HTTPException(404, f"doc not found: {filename!r}") + from core.kb_lock import KbBusyError + try: + if not delete_doc(resolve_workspace(None), user_id, name, filename): + raise HTTPException(404, f"doc not found: {filename!r}") + except KbBusyError: + raise HTTPException(409, "该库正在入库或被其他操作修改") return {"deleted": filename}