增强知识库文件写入一致性
This commit is contained in:
parent
ade3a847f6
commit
70683b7d09
|
|
@ -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 或结构设计语境。
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<hash>.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"}`。
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
97
core/kb.py
97
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 就近写进
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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/<name>`` 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
|
||||
|
|
@ -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。
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
27
tools/fs.py
27
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]"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
Loading…
Reference in New Issue