Update: 将子项目从 submodule 转为完整内容
- 移除 GovAI, nomifun-tauri, 算力盒子 的 submodule 引用 - 添加所有子项目的完整源代码 - 保留原始 .git 为 .git.bak 备份
This commit is contained in:
@@ -0,0 +1,38 @@
|
||||
[package]
|
||||
name = "nomifun-knowledge"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
|
||||
[dependencies]
|
||||
nomifun-common.workspace = true
|
||||
nomifun-db.workspace = true
|
||||
nomifun-api-types.workspace = true
|
||||
nomifun-realtime.workspace = true
|
||||
nomifun-auth.workspace = true
|
||||
nomifun-net.workspace = true
|
||||
axum.workspace = true
|
||||
tokio.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
async-trait.workspace = true
|
||||
tracing.workspace = true
|
||||
walkdir.workspace = true
|
||||
similar.workspace = true
|
||||
zip.workspace = true
|
||||
reqwest.workspace = true
|
||||
url.workspace = true
|
||||
chrono.workspace = true
|
||||
futures-util.workspace = true
|
||||
htmd.workspace = true
|
||||
base64.workspace = true
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
# NTFS junctions mount knowledge bases into workspaces without requiring
|
||||
# the SeCreateSymbolicLink privilege (same rationale as nomifun-extension).
|
||||
junction = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile.workspace = true
|
||||
wiremock.workspace = true
|
||||
tower = { workspace = true, features = ["util"] }
|
||||
http-body-util.workspace = true
|
||||
@@ -0,0 +1,403 @@
|
||||
//! Creation-time AI autogen: sample a base's markdown corpus, ask the LLM
|
||||
//! for a registry description + root README, and parse the strict-JSON
|
||||
//! reply.
|
||||
//!
|
||||
//! This module only owns the LLM **seam** ([`KnowledgeCompleter`]), the
|
||||
//! sampling/prompt/parse pure logic, and the constants. Orchestration
|
||||
//! (loading the base row, writing files, emitting events) lives in
|
||||
//! `service::KnowledgeService::generate_overview`. The production completer
|
||||
//! implementation lives in `nomifun-ai-agent` (same layering as the companion
|
||||
//! learner's `LiveCompanionCompleter`) and is late-wired via
|
||||
//! `KnowledgeService::set_completer`.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use nomifun_common::AppError;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::service::{KB_INBOX_REL_DIR, is_md};
|
||||
|
||||
/// LLM seam for knowledge autogen (same pattern as `CompanionCompleter` in
|
||||
/// `nomifun-companion`). The knowledge crate holds only the trait; provider/model
|
||||
/// selection is the implementor's concern.
|
||||
#[async_trait::async_trait]
|
||||
pub trait KnowledgeCompleter: Send + Sync {
|
||||
/// Run a one-shot completion using the implementor's default
|
||||
/// provider/model selection (the first enabled provider/model).
|
||||
async fn complete(&self, system: &str, user: &str) -> Result<String, AppError>;
|
||||
|
||||
/// Run a one-shot completion against an explicitly chosen
|
||||
/// `(provider_id, model)` instead of the default. The base trait falls
|
||||
/// back to [`Self::complete`] so existing implementations (and test
|
||||
/// fakes) keep compiling and behaving unchanged; the production
|
||||
/// completer overrides this to honor the caller's pick. Used by the
|
||||
/// user-facing autogen/description endpoints where the UI lets the user
|
||||
/// pick a model; background best-effort call sites keep using
|
||||
/// [`Self::complete`] (or pass `None`) so a transient UI choice never
|
||||
/// leaks into server-driven curation tasks.
|
||||
async fn complete_with(
|
||||
&self,
|
||||
system: &str,
|
||||
user: &str,
|
||||
_provider_id: &str,
|
||||
_model: &str,
|
||||
) -> Result<String, AppError> {
|
||||
self.complete(system, user).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Sampling budget: at most this many files feed the overview prompt.
|
||||
pub const SAMPLE_MAX_FILES: usize = 20;
|
||||
/// Sampling budget: at most this many bytes are read from each file.
|
||||
pub const SAMPLE_MAX_PER_FILE: usize = 4 * 1024;
|
||||
/// Sampling budget: total cap across all sampled excerpts.
|
||||
pub const SAMPLE_MAX_TOTAL: usize = 60 * 1024;
|
||||
|
||||
/// Generated descriptions are clamped to this many chars before persisting.
|
||||
pub const DESCRIPTION_MAX_CHARS: usize = 120;
|
||||
|
||||
/// Fetched snapshots above this size are condensed via the completer
|
||||
/// (when available) before persisting.
|
||||
pub const SNAPSHOT_COMPRESS_THRESHOLD: usize = 32 * 1024;
|
||||
/// Cap on the markdown fed into a snapshot-compression call.
|
||||
pub const SNAPSHOT_LLM_INPUT_MAX: usize = 64 * 1024;
|
||||
|
||||
/// System prompt for condensing an oversized fetched page. Output is plain
|
||||
/// markdown (not JSON).
|
||||
pub const SNAPSHOT_COMPRESS_SYSTEM: &str = "You are a knowledge-base curator. The user message is a markdown \
|
||||
document fetched from a web page; it is too long to store verbatim. Rewrite it as a condensed digest:\n\
|
||||
- Keep the original heading structure (#/##/###) where it carries meaning.\n\
|
||||
- Keep key facts, definitions, API signatures, tables and short code snippets; drop navigation, boilerplate \
|
||||
and repetition.\n\
|
||||
- Write in the document's own language.\n\
|
||||
Output ONLY the condensed markdown — no commentary, no fences.";
|
||||
|
||||
/// Strict-JSON contract for the overview generation call. Agent-facing
|
||||
/// wording is English by project convention.
|
||||
pub const OVERVIEW_SYSTEM: &str = "You are a knowledge-base curator. You will receive samples from a \
|
||||
markdown knowledge base. Reply with ONLY a JSON object of this exact shape:\n\
|
||||
{\"description\": \"...\", \"readme_markdown\": \"...\"}\n\
|
||||
Rules:\n\
|
||||
- description: one or two sentences (max 120 characters) stating what the base covers and when to \
|
||||
consult it. Write it in the dominant language of the sampled content.\n\
|
||||
- readme_markdown: a complete README.md for the base root — an H1 title, a short overview paragraph, \
|
||||
and a section describing the main topics/structure so a reader can navigate the documents. Keep the \
|
||||
README under ~300 lines; prefer a concise overview to exhaustive listings.\n\
|
||||
- Ground everything in the samples; never invent documents or facts that are not present.\n\
|
||||
- Output the JSON object only: no prose, no markdown fences.";
|
||||
|
||||
/// Strict-JSON contract for the stateless description-generation call
|
||||
/// (create-base form, before any row exists). Description only — no README.
|
||||
/// The output lands in conversation/terminal prompt contexts as
|
||||
/// `- Description: ...` under a `### {name}` heading (see `context.rs`), so
|
||||
/// it must double as a retrieval hint.
|
||||
pub const DESCRIPTION_SYSTEM: &str = "You are a knowledge-base curator. You will receive samples from a \
|
||||
markdown knowledge base. Reply with ONLY a JSON object of this exact shape:\n\
|
||||
{\"description\": \"...\"}\n\
|
||||
Rules:\n\
|
||||
- description: one or two sentences (max 120 characters) stating what topics/content the base covers \
|
||||
AND when an assistant should consult it, so a model scanning a list of bases can decide at a glance \
|
||||
whether to search this one.\n\
|
||||
- Write it in the dominant language of the sampled content.\n\
|
||||
- Ground it in the samples; never invent topics or facts that are not present.\n\
|
||||
- Output the JSON object only: no prose, no markdown fences.";
|
||||
|
||||
/// Strict-JSON contract for the stateless description-polish call: rewrite a
|
||||
/// user-typed draft into a high-quality registry description. Same prompt
|
||||
/// surface as [`DESCRIPTION_SYSTEM`] (see `context.rs` rendering).
|
||||
pub const POLISH_SYSTEM: &str = "You are a knowledge-base curator. You will receive a user-written draft \
|
||||
description of a knowledge base. Rewrite and polish the draft into a high-quality description. Reply \
|
||||
with ONLY a JSON object of this exact shape:\n\
|
||||
{\"description\": \"...\"}\n\
|
||||
Rules:\n\
|
||||
- description: one or two sentences (max 120 characters) stating what topics/content the base covers \
|
||||
AND when an assistant should consult it, so a model scanning a list of bases can decide at a glance \
|
||||
whether to search this one.\n\
|
||||
- Preserve every fact and the intent of the draft; never invent capabilities, topics or facts the \
|
||||
draft does not mention. This is a rewrite/polish, not free creation.\n\
|
||||
- Write it in the dominant language of the draft.\n\
|
||||
- Output the JSON object only: no prose, no markdown fences.";
|
||||
|
||||
/// Parsed model reply for the overview call.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
pub struct OverviewOutput {
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
#[serde(default)]
|
||||
pub readme_markdown: String,
|
||||
}
|
||||
|
||||
/// Extract the outermost `{...}` block from a raw model reply — the shared
|
||||
/// tolerance step of both parsers below (strips ```json fences and
|
||||
/// surrounding prose; same approach as the companion learner's parser).
|
||||
fn extract_json_block(raw: &str) -> Result<&str, String> {
|
||||
let start = raw.find('{').ok_or_else(|| "no JSON object found in model output".to_owned())?;
|
||||
let end = raw.rfind('}').filter(|e| *e > start).ok_or_else(|| "no JSON object found in model output".to_owned())?;
|
||||
Ok(&raw[start..=end])
|
||||
}
|
||||
|
||||
/// Parse the model output into [`OverviewOutput`], tolerating ```json fences
|
||||
/// and surrounding prose (extracts the outermost `{...}` block — same
|
||||
/// tolerance as the companion learner's parser).
|
||||
pub fn parse_overview_output(raw: &str) -> Result<OverviewOutput, String> {
|
||||
let output: OverviewOutput =
|
||||
serde_json::from_str(extract_json_block(raw)?).map_err(|e| format!("invalid overview JSON: {e}"))?;
|
||||
if output.description.trim().is_empty() && output.readme_markdown.trim().is_empty() {
|
||||
return Err("overview JSON carries neither description nor readme_markdown".into());
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Parse a description-only reply (`{"description": "..."}`) with the same
|
||||
/// tolerance as [`parse_overview_output`]: ```json fences and surrounding
|
||||
/// prose are stripped by extracting the outermost `{...}` block. An empty
|
||||
/// description is a parse failure (callers retry once).
|
||||
pub fn parse_description_output(raw: &str) -> Result<String, String> {
|
||||
#[derive(Deserialize)]
|
||||
struct DescriptionOutput {
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
}
|
||||
let output: DescriptionOutput =
|
||||
serde_json::from_str(extract_json_block(raw)?).map_err(|e| format!("invalid description JSON: {e}"))?;
|
||||
let description = output.description.trim();
|
||||
if description.is_empty() {
|
||||
return Err("description JSON carries an empty description".into());
|
||||
}
|
||||
Ok(description.to_owned())
|
||||
}
|
||||
|
||||
/// Clamp a model-produced description to [`DESCRIPTION_MAX_CHARS`] — the same
|
||||
/// bound the overview path enforces before persisting (char-based, so a
|
||||
/// multi-byte boundary can never split).
|
||||
pub fn clamp_description(raw: &str) -> String {
|
||||
raw.trim().chars().take(DESCRIPTION_MAX_CHARS).collect()
|
||||
}
|
||||
|
||||
/// Build the user prompt from the base registry info and sampled excerpts.
|
||||
pub fn build_overview_prompt(name: &str, description: &str, samples: &[(String, String)]) -> String {
|
||||
let mut prompt = format!("Knowledge base name: {name}\n");
|
||||
let description = description.trim();
|
||||
if !description.is_empty() {
|
||||
prompt.push_str(&format!("Current description: {description}\n"));
|
||||
}
|
||||
prompt.push_str(&format!("Sampled documents ({}):\n", samples.len()));
|
||||
for (rel, excerpt) in samples {
|
||||
prompt.push_str(&format!("\n--- FILE: {rel} ---\n{excerpt}\n"));
|
||||
}
|
||||
prompt.push_str("\nReply with the JSON object now.");
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Build the user prompt for the stateless description-generation call.
|
||||
/// `name` may be blank (the create form lets users ask before naming).
|
||||
pub fn build_description_prompt(name: &str, samples: &[(String, String)]) -> String {
|
||||
let mut prompt = String::new();
|
||||
let name = name.trim();
|
||||
if !name.is_empty() {
|
||||
prompt.push_str(&format!("Knowledge base name: {name}\n"));
|
||||
}
|
||||
prompt.push_str(&format!("Sampled documents ({}):\n", samples.len()));
|
||||
for (rel, excerpt) in samples {
|
||||
prompt.push_str(&format!("\n--- FILE: {rel} ---\n{excerpt}\n"));
|
||||
}
|
||||
prompt.push_str("\nReply with the JSON object now.");
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Build the user prompt for the stateless description-polish call.
|
||||
/// `name` may be blank; `draft` is the user's raw description text.
|
||||
pub fn build_polish_prompt(name: &str, draft: &str) -> String {
|
||||
let mut prompt = String::new();
|
||||
let name = name.trim();
|
||||
if !name.is_empty() {
|
||||
prompt.push_str(&format!("Knowledge base name: {name}\n"));
|
||||
}
|
||||
prompt.push_str(&format!("Draft description:\n{}\n", draft.trim()));
|
||||
prompt.push_str("\nReply with the JSON object now.");
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Sample the markdown corpus under `root` for the overview prompt:
|
||||
/// `_inbox/` (unreviewed staged write-backs) and the root `README.md` (the
|
||||
/// artifact being regenerated) are excluded; files are taken in sorted-path
|
||||
/// order up to [`SAMPLE_MAX_FILES`], each excerpt capped at
|
||||
/// [`SAMPLE_MAX_PER_FILE`] bytes, total capped at [`SAMPLE_MAX_TOTAL`].
|
||||
pub async fn sample_base_files(root: &Path) -> Vec<(String, String)> {
|
||||
let root = root.to_path_buf();
|
||||
tokio::task::spawn_blocking(move || sample_base_files_blocking(&root))
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn sample_base_files_blocking(root: &Path) -> Vec<(String, String)> {
|
||||
if !root.is_dir() {
|
||||
return Vec::new();
|
||||
}
|
||||
let mut rels: Vec<String> = walkdir::WalkDir::new(root)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|e| e.file_type().is_file() && is_md(e.path()))
|
||||
.filter_map(|e| {
|
||||
let rel = e.path().strip_prefix(root).ok()?.to_string_lossy().replace('\\', "/");
|
||||
let keep = !rel.starts_with(&format!("{KB_INBOX_REL_DIR}/")) && rel != "README.md";
|
||||
keep.then_some(rel)
|
||||
})
|
||||
.collect();
|
||||
rels.sort();
|
||||
|
||||
let mut samples = Vec::new();
|
||||
let mut total = 0usize;
|
||||
for rel in rels.into_iter().take(SAMPLE_MAX_FILES) {
|
||||
if total >= SAMPLE_MAX_TOTAL {
|
||||
break;
|
||||
}
|
||||
let budget = SAMPLE_MAX_PER_FILE.min(SAMPLE_MAX_TOTAL - total);
|
||||
let Some(excerpt) = read_prefix_lossy(&root.join(&rel), budget) else {
|
||||
continue;
|
||||
};
|
||||
if excerpt.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
total += excerpt.len();
|
||||
samples.push((rel, excerpt));
|
||||
}
|
||||
samples
|
||||
}
|
||||
|
||||
/// Read at most `limit` bytes from the start of `path`, lossily decoded
|
||||
/// (a multi-byte char cut at the boundary degrades to U+FFFD, never a panic).
|
||||
fn read_prefix_lossy(path: &Path, limit: usize) -> Option<String> {
|
||||
use std::io::Read;
|
||||
let mut file = std::fs::File::open(path).ok()?;
|
||||
let mut buf = vec![0u8; limit];
|
||||
let mut read = 0usize;
|
||||
loop {
|
||||
match file.read(&mut buf[read..]) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => read += n,
|
||||
Err(_) => return None,
|
||||
}
|
||||
if read == buf.len() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(String::from_utf8_lossy(&buf[..read]).into_owned())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_tolerates_fences_and_prose() {
|
||||
let plain = r##"{"description":"覆盖部署与运维。","readme_markdown":"# 运维库\n\n概览。"}"##;
|
||||
let out = parse_overview_output(plain).unwrap();
|
||||
assert_eq!(out.description, "覆盖部署与运维。");
|
||||
assert!(out.readme_markdown.starts_with("# 运维库"));
|
||||
|
||||
let fenced = format!("Sure, here you go:\n```json\n{plain}\n```\ndone");
|
||||
let out = parse_overview_output(&fenced).unwrap();
|
||||
assert_eq!(out.description, "覆盖部署与运维。");
|
||||
|
||||
assert!(parse_overview_output("I cannot do that").is_err());
|
||||
assert!(parse_overview_output(r#"{"description":"","readme_markdown":""}"#).is_err());
|
||||
// Missing fields default to empty (partial output still usable).
|
||||
let only_desc = parse_overview_output(r#"{"description":"d"}"#).unwrap();
|
||||
assert_eq!(only_desc.readme_markdown, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sampling_skips_inbox_and_readme_and_caps_budgets() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let root = dir.path();
|
||||
std::fs::write(root.join("README.md"), "# old readme").unwrap();
|
||||
std::fs::create_dir_all(root.join("_inbox/conv_1")).unwrap();
|
||||
std::fs::write(root.join("_inbox/conv_1/draft.md"), "# draft").unwrap();
|
||||
// 25 real files, one larger than the per-file cap.
|
||||
for i in 0..25 {
|
||||
std::fs::write(root.join(format!("f{i:02}.md")), format!("# 文件 {i}\n正文")).unwrap();
|
||||
}
|
||||
std::fs::write(root.join("big.md"), "x".repeat(SAMPLE_MAX_PER_FILE * 2)).unwrap();
|
||||
|
||||
let samples = sample_base_files(root).await;
|
||||
assert_eq!(samples.len(), SAMPLE_MAX_FILES, "{:?}", samples.iter().map(|s| &s.0).collect::<Vec<_>>());
|
||||
assert!(samples.iter().all(|(rel, _)| rel != "README.md" && !rel.starts_with("_inbox/")));
|
||||
let big = samples.iter().find(|(rel, _)| rel == "big.md").expect("big.md sampled (sorted first)");
|
||||
assert!(big.1.len() <= SAMPLE_MAX_PER_FILE);
|
||||
let total: usize = samples.iter().map(|(_, s)| s.len()).sum();
|
||||
assert!(total <= SAMPLE_MAX_TOTAL);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sampling_empty_or_missing_root() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
assert!(sample_base_files(dir.path()).await.is_empty());
|
||||
assert!(sample_base_files(&dir.path().join("nope")).await.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_carries_name_and_samples() {
|
||||
let samples = vec![("a.md".to_string(), "# A\nbody".to_string())];
|
||||
let prompt = build_overview_prompt("领域知识", "旧描述", &samples);
|
||||
assert!(prompt.contains("领域知识"));
|
||||
assert!(prompt.contains("旧描述"));
|
||||
assert!(prompt.contains("--- FILE: a.md ---"));
|
||||
assert!(prompt.contains("# A\nbody"));
|
||||
// Empty current description line is omitted.
|
||||
let prompt = build_overview_prompt("x", " ", &samples);
|
||||
assert!(!prompt.contains("Current description"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_description_tolerates_fences_and_prose() {
|
||||
let plain = r#"{"description":"覆盖部署与运维,排障时查阅。"}"#;
|
||||
assert_eq!(parse_description_output(plain).unwrap(), "覆盖部署与运维,排障时查阅。");
|
||||
|
||||
let fenced = format!("Sure!\n```json\n{plain}\n```\nthat's it");
|
||||
assert_eq!(parse_description_output(&fenced).unwrap(), "覆盖部署与运维,排障时查阅。");
|
||||
|
||||
// Whitespace-padded description is trimmed.
|
||||
assert_eq!(parse_description_output(r#"{"description":" d "}"#).unwrap(), "d");
|
||||
|
||||
// No JSON / invalid JSON / empty or missing description all fail.
|
||||
assert!(parse_description_output("I cannot do that").is_err());
|
||||
assert!(parse_description_output("{not json}").is_err());
|
||||
assert!(parse_description_output(r#"{"description":""}"#).is_err());
|
||||
assert!(parse_description_output(r#"{"other":"x"}"#).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clamp_description_caps_chars_and_trims() {
|
||||
assert_eq!(clamp_description(" short "), "short");
|
||||
let long = "知".repeat(DESCRIPTION_MAX_CHARS + 80);
|
||||
let clamped = clamp_description(&long);
|
||||
assert_eq!(clamped.chars().count(), DESCRIPTION_MAX_CHARS);
|
||||
// Exactly at the cap → kept whole.
|
||||
let exact = "k".repeat(DESCRIPTION_MAX_CHARS);
|
||||
assert_eq!(clamp_description(&exact), exact);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn description_prompt_carries_optional_name_and_samples() {
|
||||
let samples = vec![("guide.md".to_string(), "# 指南\n正文".to_string())];
|
||||
let prompt = build_description_prompt("运维库", &samples);
|
||||
assert!(prompt.contains("Knowledge base name: 运维库"));
|
||||
assert!(prompt.contains("--- FILE: guide.md ---"));
|
||||
assert!(prompt.contains("# 指南\n正文"));
|
||||
// Blank name → the name line is omitted entirely.
|
||||
let prompt = build_description_prompt(" ", &samples);
|
||||
assert!(!prompt.contains("Knowledge base name"));
|
||||
assert!(prompt.contains("Sampled documents (1):"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn polish_prompt_carries_optional_name_and_draft() {
|
||||
let prompt = build_polish_prompt("运维库", " 记录一些部署的东西 ");
|
||||
assert!(prompt.contains("Knowledge base name: 运维库"));
|
||||
assert!(prompt.contains("Draft description:\n记录一些部署的东西\n"));
|
||||
let prompt = build_polish_prompt("", "draft text");
|
||||
assert!(!prompt.contains("Knowledge base name"));
|
||||
assert!(prompt.contains("draft text"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
//! Source-connector framework. A connector pulls remote documents (Feishu
|
||||
//! wiki, Notion, …) into a managed knowledge base's `snapshots/` dir as
|
||||
//! markdown — the **"snapshot-as-seam"** invariant: connectors only produce the
|
||||
//! same markdown-file shape the URL source already does, so retrieval / mount /
|
||||
//! search / TOC stay untouched.
|
||||
//!
|
||||
//! The trait is intentionally read-oriented for v1: `push_document` is reserved
|
||||
//! (default `Err`) for future bidirectional sync but not implemented. Webhook
|
||||
//! subscription is optional (default no-op); the baseline is poll-based
|
||||
//! `list_documents` + `fetch_document`.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use nomifun_common::AppError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// A decrypted connector credential, ready to authenticate against the remote.
|
||||
/// `payload` is the connector-specific JSON (e.g. `{ "app_id", "app_secret" }`
|
||||
/// for Feishu) decrypted by the service layer from `connector_credentials`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConnectorCredential {
|
||||
pub id: String,
|
||||
pub kind: String,
|
||||
pub name: String,
|
||||
pub payload: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Identity returned by a successful credential validation (for the UI's
|
||||
/// "test connection" affordance).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ConnectorIdentity {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tenant_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub scopes_available: Vec<String>,
|
||||
}
|
||||
|
||||
/// Connector-specific scope of what to sync, e.g. a Feishu wiki space
|
||||
/// (`{ "type": "wiki_space", "space_id": "..." }`) or a Notion page list.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ConnectorScope(pub serde_json::Value);
|
||||
|
||||
/// Incremental-sync cursor. `last_sync_at` drives modified-since filtering;
|
||||
/// `opaque` carries connector-specific paging/state across runs.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SyncCursor {
|
||||
#[serde(default)]
|
||||
pub last_sync_at: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub opaque: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A reference to a remote document discovered by `list_documents` (metadata
|
||||
/// only; the body is fetched lazily by `fetch_document`).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RemoteDocRef {
|
||||
pub remote_id: String,
|
||||
pub title: String,
|
||||
/// Last-edit time (epoch ms) for incremental filtering.
|
||||
pub edit_time: i64,
|
||||
pub doc_type: String,
|
||||
}
|
||||
|
||||
/// One page of a paginated `list_documents` call.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SyncPage {
|
||||
pub docs: Vec<RemoteDocRef>,
|
||||
/// Remote ids that disappeared since the cursor (moved to `_trash/`).
|
||||
pub deleted_ids: Vec<String>,
|
||||
pub next_page_token: Option<String>,
|
||||
pub updated_cursor: SyncCursor,
|
||||
}
|
||||
|
||||
/// A fetched remote document converted to markdown, ready to snapshot.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FetchedConnectorDoc {
|
||||
pub remote_id: String,
|
||||
pub title: String,
|
||||
pub markdown: String,
|
||||
pub edit_time: i64,
|
||||
/// Canonical web URL for the snapshot frontmatter (if any).
|
||||
pub source_url: Option<String>,
|
||||
}
|
||||
|
||||
/// Result of registering a webhook (push-based sync). Optional capability.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WebhookSubscription {
|
||||
pub subscription_id: String,
|
||||
pub expires_at: Option<i64>,
|
||||
}
|
||||
|
||||
/// Reserved for future bidirectional sync (`push_document`).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PushDocumentRequest {
|
||||
pub remote_id: String,
|
||||
pub markdown: String,
|
||||
}
|
||||
|
||||
/// A pluggable source connector. Implementors own the remote API + format
|
||||
/// conversion; the sync orchestrator drives them and writes snapshots.
|
||||
#[async_trait]
|
||||
pub trait KnowledgeConnector: Send + Sync {
|
||||
/// Discriminator stored in `extra.source.kind` (e.g. "feishu").
|
||||
fn kind(&self) -> &'static str;
|
||||
|
||||
/// Validate credentials; returns connector identity. Used at credential
|
||||
/// registration time and the UI "test connection" action.
|
||||
async fn validate_credentials(&self, cred: &ConnectorCredential) -> Result<ConnectorIdentity, AppError>;
|
||||
|
||||
/// Enumerate documents in `scope`, paginated. With a populated `cursor`,
|
||||
/// returns only docs changed since `last_sync_at` (incremental); empty
|
||||
/// cursor = full sync. Removed docs are reported in `deleted_ids`.
|
||||
async fn list_documents(
|
||||
&self,
|
||||
cred: &ConnectorCredential,
|
||||
scope: &ConnectorScope,
|
||||
cursor: &SyncCursor,
|
||||
page_token: Option<&str>,
|
||||
) -> Result<SyncPage, AppError>;
|
||||
|
||||
/// Fetch one document and convert it to markdown (connector owns the
|
||||
/// format conversion, e.g. Feishu blocks → md via `feishu_md`).
|
||||
async fn fetch_document(&self, cred: &ConnectorCredential, doc: &RemoteDocRef) -> Result<FetchedConnectorDoc, AppError>;
|
||||
|
||||
/// Optional push-based sync. Default: poll-only (no webhook).
|
||||
async fn subscribe_webhook(
|
||||
&self,
|
||||
_cred: &ConnectorCredential,
|
||||
_scope: &ConnectorScope,
|
||||
_callback_url: &str,
|
||||
) -> Result<Option<WebhookSubscription>, AppError> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Reserved for future bidirectional sync. Default: unsupported.
|
||||
async fn push_document(&self, _cred: &ConnectorCredential, _doc: &PushDocumentRequest) -> Result<(), AppError> {
|
||||
Err(AppError::BadRequest("push_document is not supported by this connector".into()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
struct MockConnector;
|
||||
|
||||
#[async_trait]
|
||||
impl KnowledgeConnector for MockConnector {
|
||||
fn kind(&self) -> &'static str {
|
||||
"mock"
|
||||
}
|
||||
async fn validate_credentials(&self, _cred: &ConnectorCredential) -> Result<ConnectorIdentity, AppError> {
|
||||
Ok(ConnectorIdentity { tenant_name: Some("Acme".into()), scopes_available: vec!["wiki".into()] })
|
||||
}
|
||||
async fn list_documents(
|
||||
&self,
|
||||
_cred: &ConnectorCredential,
|
||||
_scope: &ConnectorScope,
|
||||
_cursor: &SyncCursor,
|
||||
_page_token: Option<&str>,
|
||||
) -> Result<SyncPage, AppError> {
|
||||
Ok(SyncPage {
|
||||
docs: vec![RemoteDocRef {
|
||||
remote_id: "d1".into(),
|
||||
title: "Doc One".into(),
|
||||
edit_time: 100,
|
||||
doc_type: "docx".into(),
|
||||
}],
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
async fn fetch_document(&self, _cred: &ConnectorCredential, doc: &RemoteDocRef) -> Result<FetchedConnectorDoc, AppError> {
|
||||
Ok(FetchedConnectorDoc {
|
||||
remote_id: doc.remote_id.clone(),
|
||||
title: doc.title.clone(),
|
||||
markdown: format!("# {}\n\nbody", doc.title),
|
||||
edit_time: doc.edit_time,
|
||||
source_url: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trait_is_object_safe_and_defaults_apply() {
|
||||
let c: Arc<dyn KnowledgeConnector> = Arc::new(MockConnector);
|
||||
assert_eq!(c.kind(), "mock");
|
||||
let cred = ConnectorCredential { id: "x".into(), kind: "mock".into(), name: "n".into(), payload: serde_json::json!({}) };
|
||||
assert_eq!(c.validate_credentials(&cred).await.unwrap().tenant_name.as_deref(), Some("Acme"));
|
||||
let page = c.list_documents(&cred, &ConnectorScope::default(), &SyncCursor::default(), None).await.unwrap();
|
||||
assert_eq!(page.docs.len(), 1);
|
||||
let doc = c.fetch_document(&cred, &page.docs[0]).await.unwrap();
|
||||
assert!(doc.markdown.contains("Doc One"));
|
||||
// Default capabilities.
|
||||
assert!(c.subscribe_webhook(&cred, &ConnectorScope::default(), "http://cb").await.unwrap().is_none());
|
||||
assert!(c.push_document(&cred, &PushDocumentRequest { remote_id: "d1".into(), markdown: "x".into() }).await.is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,905 @@
|
||||
//! Feishu (Lark) knowledge connector — syncs documents from a Feishu wiki space
|
||||
//! into a managed knowledge base using Feishu's Open API with a self-built app
|
||||
//! (tenant_access_token, no OAuth redirect).
|
||||
//!
|
||||
//! # Architecture
|
||||
//! - Token caching with 30-minute safety margin, automatic refresh on 401.
|
||||
//! - Rate limiting: minimum 250ms between requests (≈4 QPS), 429 retry with
|
||||
//! `x-ogw-ratelimit-reset` header.
|
||||
//! - Configurable `base_url` for wiremock test injection.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use nomifun_common::AppError;
|
||||
use reqwest::Client;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::connector::{
|
||||
ConnectorCredential, ConnectorIdentity, ConnectorScope, FetchedConnectorDoc,
|
||||
KnowledgeConnector, RemoteDocRef, SyncCursor, SyncPage,
|
||||
};
|
||||
use crate::feishu_md::blocks_to_markdown;
|
||||
|
||||
/// Safety margin subtracted from the token's stated expiry to avoid using a
|
||||
/// nearly-expired token.
|
||||
const TOKEN_SAFETY_MARGIN: Duration = Duration::from_secs(30 * 60);
|
||||
|
||||
/// Minimum interval between outgoing HTTP requests (rate limit).
|
||||
const MIN_REQUEST_INTERVAL: Duration = Duration::from_millis(250);
|
||||
|
||||
// ─── Token cache ────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct CachedToken {
|
||||
token: String,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
// ─── Internal API response structures ───────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct FeishuEnvelope<T> {
|
||||
code: i64,
|
||||
msg: String,
|
||||
data: Option<T>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TokenData {
|
||||
tenant_access_token: String,
|
||||
expire: i64, // seconds
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WikiNodesData {
|
||||
items: Option<Vec<WikiNode>>,
|
||||
has_more: Option<bool>,
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[allow(dead_code)]
|
||||
struct WikiNode {
|
||||
node_token: Option<String>,
|
||||
obj_token: Option<String>,
|
||||
obj_type: Option<String>,
|
||||
title: Option<String>,
|
||||
obj_edit_time: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct DocBlocksData {
|
||||
items: Option<Vec<Value>>,
|
||||
has_more: Option<bool>,
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
// ─── Connector ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Feishu knowledge connector.
|
||||
pub struct FeishuConnector {
|
||||
base_url: String,
|
||||
client: Client,
|
||||
token_cache: Arc<Mutex<Option<CachedToken>>>,
|
||||
last_request: Arc<Mutex<Option<Instant>>>,
|
||||
}
|
||||
|
||||
impl FeishuConnector {
|
||||
/// Create a connector with the default Feishu base URL.
|
||||
pub fn new() -> Self {
|
||||
Self::with_base_url("https://open.feishu.cn".to_string())
|
||||
}
|
||||
|
||||
/// Create a connector with a custom base URL (for testing with wiremock).
|
||||
pub fn with_base_url(base_url: String) -> Self {
|
||||
Self {
|
||||
base_url,
|
||||
client: Client::new(),
|
||||
token_cache: Arc::new(Mutex::new(None)),
|
||||
last_request: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract app_id and app_secret from credential payload.
|
||||
fn parse_credential(credential: &ConnectorCredential) -> Result<(String, String), AppError> {
|
||||
let app_id = credential
|
||||
.payload
|
||||
.get("app_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("missing app_id in credential payload".into()))?
|
||||
.to_string();
|
||||
let app_secret = credential
|
||||
.payload
|
||||
.get("app_secret")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
AppError::BadRequest("missing app_secret in credential payload".into())
|
||||
})?
|
||||
.to_string();
|
||||
Ok((app_id, app_secret))
|
||||
}
|
||||
|
||||
/// Extract space_id from scope.
|
||||
fn parse_scope(scope: &ConnectorScope) -> Result<String, AppError> {
|
||||
scope
|
||||
.0
|
||||
.get("space_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| AppError::BadRequest("missing space_id in scope".into()))
|
||||
}
|
||||
|
||||
/// Acquire a valid tenant_access_token, fetching a new one if needed.
|
||||
async fn get_token(&self, credential: &ConnectorCredential) -> Result<String, AppError> {
|
||||
// Check cache first
|
||||
{
|
||||
let cache = self.token_cache.lock().await;
|
||||
if let Some(ref cached) = *cache {
|
||||
if Instant::now() < cached.expires_at {
|
||||
return Ok(cached.token.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
// Fetch new token
|
||||
self.fetch_token(credential).await
|
||||
}
|
||||
|
||||
/// Fetch a new tenant_access_token from Feishu and cache it.
|
||||
async fn fetch_token(&self, credential: &ConnectorCredential) -> Result<String, AppError> {
|
||||
let (app_id, app_secret) = Self::parse_credential(credential)?;
|
||||
|
||||
self.rate_limit().await;
|
||||
|
||||
let url = format!(
|
||||
"{}/open-apis/auth/v3/tenant_access_token/internal",
|
||||
self.base_url
|
||||
);
|
||||
let body = serde_json::json!({
|
||||
"app_id": app_id,
|
||||
"app_secret": app_secret,
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(&url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu token request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu token read body: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu token endpoint returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let envelope: FeishuEnvelope<TokenData> = serde_json::from_str(&text).map_err(|e| {
|
||||
AppError::BadGateway(format!("feishu token parse error: {e}, body: {text}"))
|
||||
})?;
|
||||
|
||||
if envelope.code != 0 {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu token error code={}, msg={}",
|
||||
envelope.code, envelope.msg
|
||||
)));
|
||||
}
|
||||
|
||||
let data = envelope
|
||||
.data
|
||||
.ok_or_else(|| AppError::BadGateway("feishu token response missing data".into()))?;
|
||||
|
||||
let expires_at = Instant::now()
|
||||
+ Duration::from_secs(data.expire.max(0) as u64)
|
||||
- TOKEN_SAFETY_MARGIN.min(Duration::from_secs(data.expire.max(0) as u64));
|
||||
|
||||
let token = data.tenant_access_token;
|
||||
|
||||
// Cache it
|
||||
{
|
||||
let mut cache = self.token_cache.lock().await;
|
||||
*cache = Some(CachedToken {
|
||||
token: token.clone(),
|
||||
expires_at,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Invalidate the cached token (used on 401 retry).
|
||||
async fn invalidate_token(&self) {
|
||||
let mut cache = self.token_cache.lock().await;
|
||||
*cache = None;
|
||||
}
|
||||
|
||||
/// Enforce minimum interval between requests.
|
||||
async fn rate_limit(&self) {
|
||||
let mut last = self.last_request.lock().await;
|
||||
if let Some(prev) = *last {
|
||||
let elapsed = prev.elapsed();
|
||||
if elapsed < MIN_REQUEST_INTERVAL {
|
||||
tokio::time::sleep(MIN_REQUEST_INTERVAL - elapsed).await;
|
||||
}
|
||||
}
|
||||
*last = Some(Instant::now());
|
||||
}
|
||||
|
||||
/// Make an authenticated GET request with token refresh on 401 and rate-limit
|
||||
/// retry on 429.
|
||||
async fn authed_get(
|
||||
&self,
|
||||
credential: &ConnectorCredential,
|
||||
url: &str,
|
||||
) -> Result<String, AppError> {
|
||||
let token = self.get_token(credential).await?;
|
||||
|
||||
self.rate_limit().await;
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.get(url)
|
||||
.header("Authorization", format!("Bearer {token}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
|
||||
// Handle 429 rate limit
|
||||
if status.as_u16() == 429 {
|
||||
let reset_secs = resp
|
||||
.headers()
|
||||
.get("x-ogw-ratelimit-reset")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(1);
|
||||
warn!("feishu 429 rate limit, sleeping {reset_secs}s");
|
||||
tokio::time::sleep(Duration::from_secs(reset_secs)).await;
|
||||
|
||||
// Retry once
|
||||
self.rate_limit().await;
|
||||
let retry_resp = self
|
||||
.client
|
||||
.get(url)
|
||||
.header("Authorization", format!("Bearer {token}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu retry failed: {e}")))?;
|
||||
let retry_status = retry_resp.status();
|
||||
let retry_text = retry_resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu retry body: {e}")))?;
|
||||
if !retry_status.is_success() {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu retry returned {retry_status}: {retry_text}"
|
||||
)));
|
||||
}
|
||||
return Ok(retry_text);
|
||||
}
|
||||
|
||||
// Handle 401 — invalidate token and retry once
|
||||
if status.as_u16() == 401 {
|
||||
self.invalidate_token().await;
|
||||
let new_token = self.fetch_token(credential).await?;
|
||||
self.rate_limit().await;
|
||||
let retry_resp = self
|
||||
.client
|
||||
.get(url)
|
||||
.header("Authorization", format!("Bearer {new_token}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu 401 retry failed: {e}")))?;
|
||||
let retry_status = retry_resp.status();
|
||||
let retry_text = retry_resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu 401 retry body: {e}")))?;
|
||||
if !retry_status.is_success() {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu 401 retry returned {retry_status}: {retry_text}"
|
||||
)));
|
||||
}
|
||||
return Ok(retry_text);
|
||||
}
|
||||
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("feishu response body: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
/// Parse a Feishu API envelope, checking `code == 0`.
|
||||
fn parse_envelope<T: serde::de::DeserializeOwned>(text: &str) -> Result<T, AppError> {
|
||||
let envelope: FeishuEnvelope<T> = serde_json::from_str(text).map_err(|e| {
|
||||
AppError::BadGateway(format!("feishu parse error: {e}, body: {text}"))
|
||||
})?;
|
||||
if envelope.code != 0 {
|
||||
return Err(AppError::BadGateway(format!(
|
||||
"feishu error code={}, msg={}",
|
||||
envelope.code, envelope.msg
|
||||
)));
|
||||
}
|
||||
envelope
|
||||
.data
|
||||
.ok_or_else(|| AppError::BadGateway("feishu response missing data field".into()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl KnowledgeConnector for FeishuConnector {
|
||||
fn kind(&self) -> &'static str {
|
||||
"feishu"
|
||||
}
|
||||
|
||||
async fn validate_credentials(
|
||||
&self,
|
||||
credential: &ConnectorCredential,
|
||||
) -> Result<ConnectorIdentity, AppError> {
|
||||
// Fetching a token proves the app_id/secret are valid.
|
||||
self.invalidate_token().await;
|
||||
self.fetch_token(credential).await?;
|
||||
Ok(ConnectorIdentity {
|
||||
tenant_name: None,
|
||||
scopes_available: vec!["wiki".into()],
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_documents(
|
||||
&self,
|
||||
credential: &ConnectorCredential,
|
||||
scope: &ConnectorScope,
|
||||
cursor: &SyncCursor,
|
||||
page_token: Option<&str>,
|
||||
) -> Result<SyncPage, AppError> {
|
||||
let space_id = Self::parse_scope(scope)?;
|
||||
|
||||
let mut url = format!(
|
||||
"{}/open-apis/wiki/v2/spaces/{}/nodes?page_size=50",
|
||||
self.base_url, space_id
|
||||
);
|
||||
if let Some(pt) = page_token {
|
||||
url.push_str(&format!("&page_token={pt}"));
|
||||
}
|
||||
|
||||
let text = self.authed_get(credential, &url).await?;
|
||||
let data: WikiNodesData = Self::parse_envelope(&text)?;
|
||||
|
||||
let items = data.items.unwrap_or_default();
|
||||
let mut docs = Vec::new();
|
||||
|
||||
for node in items {
|
||||
// Only keep docx nodes
|
||||
let obj_type = node.obj_type.as_deref().unwrap_or("");
|
||||
if obj_type != "docx" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let obj_token = match node.obj_token {
|
||||
Some(ref t) if !t.is_empty() => t.clone(),
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let title = node.title.unwrap_or_default();
|
||||
let edit_time: i64 = node
|
||||
.obj_edit_time
|
||||
.as_deref()
|
||||
.unwrap_or("0")
|
||||
.parse()
|
||||
.unwrap_or(0);
|
||||
|
||||
// Incremental filtering: skip docs older than last_sync_at
|
||||
if let Some(last_sync) = cursor.last_sync_at {
|
||||
if edit_time <= last_sync {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
docs.push(RemoteDocRef {
|
||||
remote_id: obj_token,
|
||||
title,
|
||||
edit_time,
|
||||
doc_type: "docx".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let next_page_token = if data.has_more.unwrap_or(false) {
|
||||
data.page_token
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// The updated cursor captures the max edit_time seen so far
|
||||
let max_edit = docs.iter().map(|d| d.edit_time).max();
|
||||
let updated_last_sync = match (cursor.last_sync_at, max_edit) {
|
||||
(Some(prev), Some(new)) => Some(prev.max(new)),
|
||||
(None, Some(new)) => Some(new),
|
||||
(prev, None) => prev,
|
||||
};
|
||||
|
||||
Ok(SyncPage {
|
||||
docs,
|
||||
deleted_ids: Vec::new(),
|
||||
next_page_token,
|
||||
updated_cursor: SyncCursor {
|
||||
last_sync_at: updated_last_sync,
|
||||
opaque: serde_json::Value::Null,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_document(
|
||||
&self,
|
||||
credential: &ConnectorCredential,
|
||||
doc: &RemoteDocRef,
|
||||
) -> Result<FetchedConnectorDoc, AppError> {
|
||||
let remote_id = doc.remote_id.as_str();
|
||||
// Paginate all blocks
|
||||
let mut all_blocks: Vec<Value> = Vec::new();
|
||||
let mut page_token: Option<String> = None;
|
||||
|
||||
loop {
|
||||
let mut url = format!(
|
||||
"{}/open-apis/docx/v1/documents/{}/blocks?page_size=500",
|
||||
self.base_url, remote_id
|
||||
);
|
||||
if let Some(ref pt) = page_token {
|
||||
url.push_str(&format!("&page_token={pt}"));
|
||||
}
|
||||
|
||||
let text = self.authed_get(credential, &url).await?;
|
||||
let data: DocBlocksData = Self::parse_envelope(&text)?;
|
||||
|
||||
if let Some(items) = data.items {
|
||||
all_blocks.extend(items);
|
||||
}
|
||||
|
||||
if data.has_more.unwrap_or(false) {
|
||||
page_token = data.page_token;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let markdown = blocks_to_markdown(&all_blocks);
|
||||
|
||||
// Best-effort source URL (uses the default feishu host for the doc link,
|
||||
// not the base_url which may be a test server).
|
||||
let source_url = format!("https://open.feishu.cn/docx/{remote_id}");
|
||||
|
||||
Ok(FetchedConnectorDoc {
|
||||
remote_id: doc.remote_id.clone(),
|
||||
title: doc.title.clone(),
|
||||
markdown,
|
||||
edit_time: doc.edit_time,
|
||||
source_url: Some(source_url),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ──────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use wiremock::matchers::{method, path, query_param};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
fn test_credential() -> ConnectorCredential {
|
||||
ConnectorCredential {
|
||||
id: "cred-1".into(),
|
||||
kind: "feishu".into(),
|
||||
name: "test".into(),
|
||||
payload: json!({
|
||||
"app_id": "cli_test123",
|
||||
"app_secret": "secret_abc"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn test_scope() -> ConnectorScope {
|
||||
ConnectorScope(json!({"space_id": "space_xyz"}))
|
||||
}
|
||||
|
||||
fn token_response() -> Value {
|
||||
json!({
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": {
|
||||
"tenant_access_token": "t-fake-token-abc",
|
||||
"expire": 7200
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn wiki_nodes_response() -> Value {
|
||||
json!({
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": {
|
||||
"items": [
|
||||
{
|
||||
"node_token": "node1",
|
||||
"obj_token": "doc_obj_1",
|
||||
"obj_type": "docx",
|
||||
"title": "Design Doc",
|
||||
"obj_edit_time": "1700000000000"
|
||||
},
|
||||
{
|
||||
"node_token": "node2",
|
||||
"obj_token": "sheet_obj_2",
|
||||
"obj_type": "sheet",
|
||||
"title": "Budget Sheet",
|
||||
"obj_edit_time": "1700000001000"
|
||||
},
|
||||
{
|
||||
"node_token": "node3",
|
||||
"obj_token": "doc_obj_3",
|
||||
"obj_type": "docx",
|
||||
"title": "API Reference",
|
||||
"obj_edit_time": "1700000002000"
|
||||
}
|
||||
],
|
||||
"has_more": false,
|
||||
"page_token": null
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn doc_blocks_response() -> Value {
|
||||
json!({
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": {
|
||||
"items": [
|
||||
{
|
||||
"block_id": "doc_root",
|
||||
"parent_id": "",
|
||||
"block_type": 1,
|
||||
"children": ["blk_h", "blk_p"]
|
||||
},
|
||||
{
|
||||
"block_id": "blk_h",
|
||||
"parent_id": "doc_root",
|
||||
"block_type": 3,
|
||||
"children": [],
|
||||
"heading1": {
|
||||
"elements": [{"text_run": {"content": "Hello Feishu"}}]
|
||||
}
|
||||
},
|
||||
{
|
||||
"block_id": "blk_p",
|
||||
"parent_id": "doc_root",
|
||||
"block_type": 2,
|
||||
"children": [],
|
||||
"text": {
|
||||
"elements": [{"text_run": {"content": "This is a test document."}}]
|
||||
}
|
||||
}
|
||||
],
|
||||
"has_more": false,
|
||||
"page_token": null
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn setup_token_mock(server: &MockServer) {
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/open-apis/auth/v3/tenant_access_token/internal"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(token_response()))
|
||||
.mount(server)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_credentials_success() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
|
||||
let identity = connector.validate_credentials(&cred).await.unwrap();
|
||||
assert!(identity.tenant_name.is_none());
|
||||
assert_eq!(identity.scopes_available, vec!["wiki"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_credentials_bad_secret() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/open-apis/auth/v3/tenant_access_token/internal"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"code": 10003,
|
||||
"msg": "app_secret is invalid",
|
||||
"data": null
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
|
||||
let err = connector.validate_credentials(&cred).await.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("10003"), "expected code in error: {msg}");
|
||||
assert!(msg.contains("app_secret is invalid"), "expected msg: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_documents_filters_non_docx() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/wiki/v2/spaces/space_xyz/nodes"))
|
||||
.and(query_param("page_size", "50"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(wiki_nodes_response()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let scope = test_scope();
|
||||
let cursor = SyncCursor::default();
|
||||
|
||||
let page = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should have 2 docx docs, not the sheet
|
||||
assert_eq!(page.docs.len(), 2);
|
||||
assert_eq!(page.docs[0].remote_id, "doc_obj_1");
|
||||
assert_eq!(page.docs[0].title, "Design Doc");
|
||||
assert_eq!(page.docs[0].edit_time, 1700000000000);
|
||||
assert_eq!(page.docs[0].doc_type, "docx");
|
||||
|
||||
assert_eq!(page.docs[1].remote_id, "doc_obj_3");
|
||||
assert_eq!(page.docs[1].title, "API Reference");
|
||||
assert_eq!(page.docs[1].edit_time, 1700000002000);
|
||||
|
||||
assert!(page.next_page_token.is_none());
|
||||
assert!(page.deleted_ids.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_documents_incremental_cursor() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/wiki/v2/spaces/space_xyz/nodes"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(wiki_nodes_response()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let scope = test_scope();
|
||||
|
||||
// Set cursor to filter out doc_obj_1 (edit_time 1700000000000)
|
||||
let cursor = SyncCursor {
|
||||
last_sync_at: Some(1700000000000),
|
||||
opaque: serde_json::Value::Null,
|
||||
};
|
||||
|
||||
let page = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Only doc_obj_3 (edit_time 1700000002000 > 1700000000000) should remain
|
||||
assert_eq!(page.docs.len(), 1);
|
||||
assert_eq!(page.docs[0].remote_id, "doc_obj_3");
|
||||
assert_eq!(page.docs[0].edit_time, 1700000002000);
|
||||
|
||||
// Updated cursor should reflect the max
|
||||
assert_eq!(page.updated_cursor.last_sync_at, Some(1700000002000));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fetch_document_converts_blocks() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/docx/v1/documents/doc_obj_1/blocks"))
|
||||
.and(query_param("page_size", "500"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(doc_blocks_response()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
|
||||
let doc = connector
|
||||
.fetch_document(
|
||||
&cred,
|
||||
&RemoteDocRef {
|
||||
remote_id: "doc_obj_1".into(),
|
||||
title: "Hello Feishu".into(),
|
||||
edit_time: 1700000000000,
|
||||
doc_type: "docx".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(doc.remote_id, "doc_obj_1");
|
||||
assert_eq!(doc.edit_time, 1700000000000);
|
||||
assert!(doc.markdown.contains("# Hello Feishu"));
|
||||
assert!(doc.markdown.contains("This is a test document."));
|
||||
assert!(doc.source_url.as_deref().unwrap_or("").contains("doc_obj_1"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_error_code_surfaces_app_error() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/wiki/v2/spaces/space_xyz/nodes"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"code": 99991,
|
||||
"msg": "internal server error from feishu",
|
||||
"data": null
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let scope = test_scope();
|
||||
let cursor = SyncCursor::default();
|
||||
|
||||
let err = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("99991"), "error should contain code: {msg}");
|
||||
assert!(
|
||||
msg.contains("internal server error from feishu"),
|
||||
"error should contain msg: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_cache_reuses_token() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Set up token mock with expect(1) — should only be called once
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/open-apis/auth/v3/tenant_access_token/internal"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(token_response()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/wiki/v2/spaces/space_xyz/nodes"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(wiki_nodes_response()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let scope = test_scope();
|
||||
let cursor = SyncCursor::default();
|
||||
|
||||
// First call — fetches token
|
||||
let _page1 = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Second call — should reuse cached token (no second POST to token endpoint)
|
||||
let _page2 = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// wiremock will verify expect(1) on drop — if token was requested twice,
|
||||
// the test will panic with "Expected exactly 1 matching request, got 2"
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pagination_has_more() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
// First page with has_more=true
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/open-apis/wiki/v2/spaces/space_xyz/nodes"))
|
||||
.and(query_param("page_size", "50"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": {
|
||||
"items": [{
|
||||
"node_token": "n1",
|
||||
"obj_token": "d1",
|
||||
"obj_type": "docx",
|
||||
"title": "Page 1 Doc",
|
||||
"obj_edit_time": "1700000000000"
|
||||
}],
|
||||
"has_more": true,
|
||||
"page_token": "next_page_abc"
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let scope = test_scope();
|
||||
let cursor = SyncCursor::default();
|
||||
|
||||
let page = connector
|
||||
.list_documents(&cred, &scope, &cursor, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(page.docs.len(), 1);
|
||||
assert_eq!(page.next_page_token, Some("next_page_abc".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kind_returns_feishu() {
|
||||
let connector = FeishuConnector::new();
|
||||
assert_eq!(connector.kind(), "feishu");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_app_id_errors() {
|
||||
let connector = FeishuConnector::new();
|
||||
let cred = ConnectorCredential {
|
||||
id: "x".into(),
|
||||
kind: "feishu".into(),
|
||||
name: "bad".into(),
|
||||
payload: json!({"app_secret": "s"}),
|
||||
};
|
||||
let err = connector.validate_credentials(&cred).await.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("app_id"), "should mention app_id: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_space_id_errors() {
|
||||
let server = MockServer::start().await;
|
||||
setup_token_mock(&server).await;
|
||||
|
||||
let connector = FeishuConnector::with_base_url(server.uri());
|
||||
let cred = test_credential();
|
||||
let bad_scope = ConnectorScope(json!({}));
|
||||
let cursor = SyncCursor::default();
|
||||
|
||||
let err = connector
|
||||
.list_documents(&cred, &bad_scope, &cursor, None)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("space_id"), "should mention space_id: {msg}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,981 @@
|
||||
//! Shared knowledge-context builder — the single source of truth for the
|
||||
//! prompt/document text that tells an agent which knowledge bases are
|
||||
//! mounted, how to retrieve from them, and which write-back contract
|
||||
//! applies.
|
||||
//!
|
||||
//! Consumers:
|
||||
//! - `nomifun-ai-agent` factory paths (ACP assembler preset context, nomi
|
||||
//! engine system prompt) via [`KnowledgeContextFormat::PromptSection`];
|
||||
//! - the terminal-session task (C1) writes a standalone
|
||||
//! `{cwd}/.nomi/knowledge/README.md` via
|
||||
//! [`KnowledgeContextFormat::TerminalReadme`].
|
||||
//!
|
||||
//! All agent-facing contract wording is English by project convention.
|
||||
|
||||
use nomifun_api_types::KnowledgeMountInfo;
|
||||
|
||||
/// Per-base cap on TOC file lines injected into the context (bounds token
|
||||
/// cost while keeping enough navigation surface for hit rate).
|
||||
pub const TOC_PER_KB_MAX: usize = 20;
|
||||
|
||||
/// Global cap on TOC file lines across all mounted bases. When many bases
|
||||
/// are mounted the per-base budget shrinks to `TOC_GLOBAL_MAX / n`.
|
||||
pub const TOC_GLOBAL_MAX: usize = 60;
|
||||
|
||||
/// Typed write-back mode, parsed ONCE at the context-builder entry point.
|
||||
/// The wire/API surfaces keep passing strings
|
||||
/// ([`KnowledgeContextOptions::writeback_mode`] stays `Option<&str>`); this
|
||||
/// enum replaces the internal string comparisons so a typo'd mode can never
|
||||
/// silently pick a branch.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum WritebackMode {
|
||||
/// Agent writes are confined to `_inbox/{target_id}/` (the safe default).
|
||||
#[default]
|
||||
Staged,
|
||||
/// The agent may edit the base body directly.
|
||||
Direct,
|
||||
}
|
||||
|
||||
impl WritebackMode {
|
||||
/// Parse a wire string: `None`/`"staged"` → [`Self::Staged`], `"direct"`
|
||||
/// → [`Self::Direct`]. Unknown values fall back to the safe default
|
||||
/// ([`Self::Staged`]) with a warning — never to the more permissive mode.
|
||||
pub fn parse(raw: Option<&str>) -> Self {
|
||||
match raw {
|
||||
None | Some("staged") => Self::Staged,
|
||||
Some("direct") => Self::Direct,
|
||||
Some(other) => {
|
||||
tracing::warn!(writeback_mode = other, "unknown writeback_mode; falling back to staged");
|
||||
Self::Staged
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Typed write-back disposition ("回写意识"), parsed ONCE at the
|
||||
/// context-builder entry point. ORTHOGONAL to [`WritebackMode`]: the mode
|
||||
/// decides WHERE writes land (staged inbox vs direct body), the eagerness
|
||||
/// decides HOW EAGERLY the agent writes at all. The wire/API surfaces keep
|
||||
/// passing strings ([`KnowledgeContextOptions::writeback_eagerness`] stays
|
||||
/// `Option<&str>`); this enum replaces internal string comparisons.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum WritebackEagerness {
|
||||
/// Restrained: only persist knowledge the model judges clearly worth
|
||||
/// keeping. The historical behaviour and the safe default.
|
||||
#[default]
|
||||
Conservative,
|
||||
/// Bold: capture anything plausibly relevant to a mounted base without
|
||||
/// much hesitation; the user prunes later.
|
||||
Aggressive,
|
||||
}
|
||||
|
||||
impl WritebackEagerness {
|
||||
/// Parse a wire string: `None`/`"conservative"` → [`Self::Conservative`],
|
||||
/// `"aggressive"` → [`Self::Aggressive`]. Unknown values fall back to the
|
||||
/// restrained default ([`Self::Conservative`]) with a warning — never to
|
||||
/// the more eager mode.
|
||||
pub fn parse(raw: Option<&str>) -> Self {
|
||||
match raw {
|
||||
None | Some("conservative") => Self::Conservative,
|
||||
Some("aggressive") => Self::Aggressive,
|
||||
Some(other) => {
|
||||
tracing::warn!(
|
||||
writeback_eagerness = other,
|
||||
"unknown writeback_eagerness; falling back to conservative"
|
||||
);
|
||||
Self::Conservative
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Output shape of [`build_knowledge_context`].
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum KnowledgeContextFormat {
|
||||
/// A `## Knowledge bases (extended knowledge source)` section meant to be
|
||||
/// embedded into a larger system prompt / preset context.
|
||||
PromptSection,
|
||||
/// A standalone, complete markdown document (H1 + intro) meant to be
|
||||
/// written as `README.md` inside the workspace mount directory.
|
||||
TerminalReadme,
|
||||
}
|
||||
|
||||
/// Inputs beyond the mounts themselves. `target_id` is the session-scoped
|
||||
/// identifier (conversation id today, terminal id for C1) used to scope the
|
||||
/// staged write-back inbox path `_inbox/{target_id}/`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct KnowledgeContextOptions<'a> {
|
||||
pub format: KnowledgeContextFormat,
|
||||
/// Write-back ("回血") switch — `false` renders the read-only contract.
|
||||
pub writeback: bool,
|
||||
/// `staged` (default when `None`) or `direct`; only meaningful while
|
||||
/// `writeback` is true.
|
||||
pub writeback_mode: Option<&'a str>,
|
||||
/// `conservative` (default when `None`) or `aggressive`; the write-back
|
||||
/// disposition ("回写意识"), only meaningful while `writeback` is true.
|
||||
pub writeback_eagerness: Option<&'a str>,
|
||||
/// Conversation / terminal id scoping staged write-backs.
|
||||
pub target_id: &'a str,
|
||||
/// Whether THIS surface exposes a `knowledge_search` agent tool. When true,
|
||||
/// the protocol leads with an imperative to call it; when false (e.g. a raw
|
||||
/// terminal PTY, or an ACP session before the knowledge MCP exists), it
|
||||
/// keeps the Grep/Read file-navigation wording.
|
||||
pub has_search_tool: bool,
|
||||
/// Whether THIS surface exposes the native `knowledge_write` agent tool.
|
||||
/// When true, the write-back contract tells the agent to CALL it — the
|
||||
/// reliable path for nomi-engine sessions, where the generic `Write` tool
|
||||
/// has no workspace cwd (relative mount paths miss the base) and sits behind
|
||||
/// the approval gate. When false (terminal PTY, ACP file-based sessions),
|
||||
/// the contract keeps the file-write prose against the mounted directory.
|
||||
pub has_write_tool: bool,
|
||||
}
|
||||
|
||||
/// Render the knowledge context for the given mounts. Returns `None` when
|
||||
/// nothing is mounted (callers skip the section entirely).
|
||||
pub fn build_knowledge_context(
|
||||
mounts: &[KnowledgeMountInfo],
|
||||
options: &KnowledgeContextOptions<'_>,
|
||||
) -> Option<String> {
|
||||
if mounts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Exhaustive on purpose: a future format variant must consciously pick
|
||||
// its rendering branch instead of silently falling into one of them.
|
||||
let readme = match options.format {
|
||||
KnowledgeContextFormat::TerminalReadme => true,
|
||||
KnowledgeContextFormat::PromptSection => false,
|
||||
};
|
||||
// Parse the writeback mode once at the entry point; everything below
|
||||
// works on the typed value.
|
||||
let writeback_mode = WritebackMode::parse(options.writeback_mode);
|
||||
let writeback_eagerness = WritebackEagerness::parse(options.writeback_eagerness);
|
||||
let mut out = String::new();
|
||||
|
||||
// ── Header ───────────────────────────────────────────────────────
|
||||
if readme {
|
||||
out.push_str(
|
||||
"# Knowledge bases\n\n\
|
||||
This directory is mounted and managed by the NomiFun platform. It contains the \
|
||||
knowledge bases bound to this session — a curated, extended knowledge source for \
|
||||
your work here. Paths below are relative to the workspace root.\n\n\
|
||||
## Retrieval protocol\n\n",
|
||||
);
|
||||
} else {
|
||||
out.push_str(
|
||||
"## Knowledge bases (extended knowledge source)\n\
|
||||
The following knowledge bases are mounted into this workspace as markdown \
|
||||
directories — a curated, extended knowledge source for this session.\n\n\
|
||||
Retrieval protocol:\n",
|
||||
);
|
||||
}
|
||||
|
||||
// ── Retrieval protocol (rendered once, not per base) ─────────────
|
||||
if options.has_search_tool {
|
||||
out.push_str(
|
||||
"1. Search first, then answer: when a task or question touches any topic covered \
|
||||
below, call the `knowledge_search` tool BEFORE answering from memory. It searches \
|
||||
the real base content directly (so it finds matches even when Grep/Glob cannot) and \
|
||||
returns ranked `base / path — heading` results, each with an opaque `handle`.\n\
|
||||
2. To read a full document, call the `knowledge_read` tool with its `handle` (no path \
|
||||
needed). The per-base tables of contents below are a map for browsing when you already \
|
||||
know the structure.\n",
|
||||
);
|
||||
} else {
|
||||
out.push_str(
|
||||
"1. Search first, then answer: when a task or question touches any topic covered \
|
||||
below, consult the matching knowledge base BEFORE answering from memory.\n\
|
||||
2. Locate documents via each base's table of contents, then read the file. For \
|
||||
anything not listed, search the base's mount path with Grep/Glob instead of \
|
||||
crawling directories blindly.\n",
|
||||
);
|
||||
}
|
||||
out.push_str(
|
||||
"3. A line like `docs/ — 12 files` summarizes a folder too large to list in full; \
|
||||
explore that folder directly when it looks relevant.\n\
|
||||
4. When you cite knowledge in an answer, reference the source file by its relative \
|
||||
path inside the mount.\n\
|
||||
5. ",
|
||||
);
|
||||
out.push_str(&writeback_contract(options, writeback_mode, writeback_eagerness));
|
||||
out.push('\n');
|
||||
|
||||
// ── Per-base sections ─────────────────────────────────────────────
|
||||
if readme {
|
||||
out.push_str("\n## Mounted bases\n");
|
||||
}
|
||||
for m in mounts {
|
||||
out.push_str(&format!("\n### {}\n", m.name));
|
||||
out.push_str(&format!("- Path: `./{}/`\n", m.rel_path));
|
||||
let description = m.description.trim();
|
||||
if !description.is_empty() {
|
||||
out.push_str(&format!("- Description: {description}\n"));
|
||||
}
|
||||
if let Some(summary) = m.summary.as_deref().map(str::trim).filter(|s| !s.is_empty()) {
|
||||
out.push_str(&format!("- Summary: {summary}\n"));
|
||||
}
|
||||
let has_summary = m.summary.as_deref().map(str::trim).is_some_and(|s| !s.is_empty());
|
||||
if description.is_empty() && !has_summary {
|
||||
let hints = toc_topic_hints(&m.toc);
|
||||
if !hints.is_empty() {
|
||||
out.push_str(&format!("- Topics include: {hints}\n"));
|
||||
}
|
||||
}
|
||||
out.push_str(&format!(
|
||||
"- When to consult: any task or question related to \"{}\" or the topics above — \
|
||||
read the matching documents first.\n",
|
||||
m.name
|
||||
));
|
||||
if !m.toc.is_empty() {
|
||||
out.push_str("- Contents:\n");
|
||||
for entry in &m.toc {
|
||||
out.push_str(&format!(" - {entry}\n"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Realtime (live URL) sources ───────────────────────────────────
|
||||
if mounts.iter().any(|m| !m.live_sources.is_empty()) {
|
||||
out.push_str(if readme { "\n## Realtime sources\n" } else { "\n### Realtime sources\n" });
|
||||
out.push_str(
|
||||
"Some bases are backed by live URL sources; their mounted snapshots may be stale. \
|
||||
When freshness matters, fetch the URL directly:\n",
|
||||
);
|
||||
for m in mounts {
|
||||
for src in &m.live_sources {
|
||||
match src.title.as_deref().map(str::trim).filter(|t| !t.is_empty()) {
|
||||
Some(title) => out.push_str(&format!("- {title} — {} (base: \"{}\")\n", src.url, m.name)),
|
||||
None => out.push_str(&format!("- {} (base: \"{}\")\n", src.url, m.name)),
|
||||
}
|
||||
}
|
||||
}
|
||||
// Layered tool guidance: `nomi_knowledge_fetch_url` is a desktop
|
||||
// gateway tool — terminal CLI sessions and plain chat sessions never
|
||||
// have it, so the text must not promise it unconditionally.
|
||||
out.push_str(
|
||||
"To read one of these URLs at its current state, use a web-fetch tool already \
|
||||
available in this session if you have one; otherwise, if a \
|
||||
`nomi_knowledge_fetch_url` tool is available, use that. If neither is \
|
||||
available, do not improvise: tell the user that this session cannot read \
|
||||
realtime sources, and answer from the mounted snapshots while noting they \
|
||||
may be stale.\n",
|
||||
);
|
||||
}
|
||||
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Up to 6 document headings pulled from a base's budgeted TOC, joined with
|
||||
/// "; ", for the "Topics include" hint shown when a base has no
|
||||
/// description/summary. Skips aggregate rows (`dir/ — N files`) and the
|
||||
/// `(+N more files)` remainder so only real document titles become hints.
|
||||
fn toc_topic_hints(toc: &[String]) -> String {
|
||||
toc.iter()
|
||||
.filter_map(|line| line.split_once(" — ").map(|(_, title)| title.trim()))
|
||||
.filter(|t| !t.is_empty() && !t.ends_with(" files") && *t != "files")
|
||||
.take(6)
|
||||
.collect::<Vec<_>>()
|
||||
.join("; ")
|
||||
}
|
||||
|
||||
/// The write-back ("回血") contract paragraph. Wording is load-bearing:
|
||||
/// staged mode confines writes to `_inbox/{target_id}/`, direct mode allows
|
||||
/// editing the base body, disabled declares everything read-only. When
|
||||
/// write-back is enabled, the disposition (`eagerness`) sentence is appended
|
||||
/// to tune HOW EAGERLY the agent writes — orthogonal to the staged/direct
|
||||
/// placement decision.
|
||||
fn writeback_contract(
|
||||
options: &KnowledgeContextOptions<'_>,
|
||||
mode: WritebackMode,
|
||||
eagerness: WritebackEagerness,
|
||||
) -> String {
|
||||
if !options.writeback {
|
||||
return "Write-back is DISABLED for this session: treat these directories as READ-ONLY. \
|
||||
Do not create, modify, or delete any files inside them."
|
||||
.to_owned();
|
||||
}
|
||||
// Tool-based contract: the surface exposes the native `knowledge_write`
|
||||
// tool, so instruct the agent to CALL it. This is the reliable path for the
|
||||
// nomi engine — the tool resolves the base + placement internally and is
|
||||
// allow-listed past the approval gate, unlike the generic Write tool.
|
||||
let mut contract = if options.has_write_tool {
|
||||
match mode {
|
||||
WritebackMode::Staged => "Write-back is ENABLED in STAGED mode: when you produce reusable knowledge \
|
||||
(conclusions, domain facts, lessons learned), persist it by CALLING the `knowledge_write` tool. \
|
||||
To UPDATE an existing document, pass the `handle` from a `knowledge_search` result (read it first \
|
||||
with `knowledge_read`, merge, then write the full `content`); to CREATE a new one, pass `base` plus \
|
||||
a descriptive `.md` `rel_path`. The system automatically places your write in a review inbox keyed \
|
||||
to this session — you do NOT manage the path, and the original document is left untouched for the \
|
||||
user to merge later. Never rebuild paths by hand. Do NOT use the generic Write/Edit file tools; \
|
||||
treat the mounted base files as READ-ONLY."
|
||||
.to_owned(),
|
||||
WritebackMode::Direct => "Write-back is ENABLED in DIRECT mode: when you produce reusable knowledge \
|
||||
(conclusions, domain facts, lessons learned), persist it by CALLING the `knowledge_write` tool — it \
|
||||
writes straight into the matching knowledge base. To UPDATE an existing document, pass the `handle` \
|
||||
from a `knowledge_search` result (read it first with `knowledge_read`, merge, then write the full \
|
||||
`content`); to CREATE a new one, pass `base` plus a descriptive `.md` `rel_path`. Never rebuild \
|
||||
paths by hand. Do NOT use the generic Write/Edit file tools for knowledge; never delete files."
|
||||
.to_owned(),
|
||||
}
|
||||
} else {
|
||||
match mode {
|
||||
WritebackMode::Staged => format!(
|
||||
"Write-back is ENABLED in STAGED mode: when you produce reusable knowledge \
|
||||
(conclusions, domain facts, lessons learned), distill it into well-structured \
|
||||
markdown files and save them ONLY under `_inbox/{}/` inside the matching \
|
||||
knowledge base directory (create it if missing). Treat everything else in the \
|
||||
knowledge bases as READ-ONLY — never modify or delete existing documents. \
|
||||
Staged notes are reviewed and merged by the user later, so make each file \
|
||||
self-contained, concise, and free of session-specific noise.",
|
||||
options.target_id
|
||||
),
|
||||
WritebackMode::Direct => {
|
||||
"Write-back is ENABLED in DIRECT mode: when you produce reusable knowledge \
|
||||
(conclusions, domain facts, lessons learned), distill it into well-structured \
|
||||
markdown files inside the matching knowledge base directory — create new files or \
|
||||
make small, focused updates to existing ones. Never rewrite documents wholesale \
|
||||
and never delete files; other sessions may be using the same base concurrently. \
|
||||
Keep entries concise, organized, and free of session-specific noise."
|
||||
.to_owned()
|
||||
}
|
||||
}
|
||||
};
|
||||
contract.push(' ');
|
||||
contract.push_str(eagerness_clause(eagerness));
|
||||
contract
|
||||
}
|
||||
|
||||
/// The write-back disposition ("回写意识") sentence appended to an enabled
|
||||
/// write-back contract. Only the threshold for WHAT to write changes — the
|
||||
/// placement rules (staged/direct) above are unaffected.
|
||||
fn eagerness_clause(eagerness: WritebackEagerness) -> &'static str {
|
||||
match eagerness {
|
||||
WritebackEagerness::Conservative => {
|
||||
"Disposition — CONSERVATIVE: be restrained about what you write back. Persist only \
|
||||
durable, broadly reusable knowledge you judge clearly worth keeping; when in doubt, \
|
||||
do NOT write. Skip session-specific, trivial, redundant, or uncertain material."
|
||||
}
|
||||
WritebackEagerness::Aggressive => {
|
||||
"Disposition — AGGRESSIVE: be eager to write back. Whenever you encounter or produce \
|
||||
anything plausibly relevant to a mounted base — facts, decisions, useful snippets, \
|
||||
observations, gotchas — capture it without much hesitation, even if you are unsure it \
|
||||
will be reused. Prefer over-capturing to losing knowledge; the user prunes later. \
|
||||
Still skip secrets and pure session noise, and keep each entry self-contained."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Max aggregated `dir/ — N files` rows appended per base; overflow beyond
|
||||
/// the largest directories folds into the `(+N more files)` row.
|
||||
pub const TOC_AGGREGATE_DIR_MAX: usize = 8;
|
||||
|
||||
/// Apply the per-KB / global TOC budgets to full per-base file listings
|
||||
/// (one inner `Vec` per mounted base, lines sorted by path).
|
||||
///
|
||||
/// The budget caps **file lines only**: a base keeps at most
|
||||
/// `min(TOC_PER_KB_MAX, TOC_GLOBAL_MAX / n_bases)` individual file rows, so
|
||||
/// file lines never exceed `TOC_GLOBAL_MAX` in total. Overflowing files are
|
||||
/// then summarized in ADDITIONAL aggregate rows — up to
|
||||
/// [`TOC_AGGREGATE_DIR_MAX`] `dir/ — N files` rows (the top-level
|
||||
/// directories with the most overflow, rendered in path order) plus one
|
||||
/// `(+N more files)` row counting root-level overflow and any directories
|
||||
/// beyond the top-8. A budgeted TOC may therefore be up to 9 rows longer
|
||||
/// than its file-line budget.
|
||||
pub fn apply_toc_budgets(tocs: &mut [Vec<String>]) {
|
||||
if tocs.is_empty() {
|
||||
return;
|
||||
}
|
||||
let per_kb = TOC_PER_KB_MAX.min((TOC_GLOBAL_MAX / tocs.len()).max(1));
|
||||
for toc in tocs.iter_mut() {
|
||||
if toc.len() <= per_kb {
|
||||
continue;
|
||||
}
|
||||
let overflow = toc.split_off(per_kb);
|
||||
let mut dirs: std::collections::BTreeMap<&str, usize> = std::collections::BTreeMap::new();
|
||||
let mut rest = 0usize;
|
||||
for line in &overflow {
|
||||
// The path is everything before the ` — title` suffix (if any).
|
||||
let path = line.split(" — ").next().unwrap_or(line);
|
||||
match path.split_once('/') {
|
||||
Some((dir, _)) => *dirs.entry(dir).or_default() += 1,
|
||||
None => rest += 1,
|
||||
}
|
||||
}
|
||||
// Keep only the heaviest directories as named rows; everything else
|
||||
// joins the `(+N more files)` remainder.
|
||||
let mut dir_counts: Vec<(&str, usize)> = dirs.into_iter().collect();
|
||||
if dir_counts.len() > TOC_AGGREGATE_DIR_MAX {
|
||||
dir_counts.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
|
||||
for (_, n) in dir_counts.split_off(TOC_AGGREGATE_DIR_MAX) {
|
||||
rest += n;
|
||||
}
|
||||
dir_counts.sort_by(|a, b| a.0.cmp(b.0));
|
||||
}
|
||||
toc.extend(dir_counts.into_iter().map(|(dir, n)| format!("{dir}/ — {n} files")));
|
||||
if rest > 0 {
|
||||
toc.push(format!("(+{rest} more files)"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use nomifun_api_types::KnowledgeSourceEntry;
|
||||
|
||||
fn mount(name: &str, rel: &str) -> KnowledgeMountInfo {
|
||||
KnowledgeMountInfo {
|
||||
id: format!("kb_{name}"),
|
||||
name: name.to_owned(),
|
||||
description: String::new(),
|
||||
rel_path: rel.to_owned(),
|
||||
toc: Vec::new(),
|
||||
summary: None,
|
||||
live_sources: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_opts<'a>(writeback: bool, mode: Option<&'a str>, target: &'a str) -> KnowledgeContextOptions<'a> {
|
||||
KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::PromptSection,
|
||||
writeback,
|
||||
writeback_mode: mode,
|
||||
writeback_eagerness: None,
|
||||
target_id: target,
|
||||
has_search_tool: false,
|
||||
has_write_tool: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Like [`prompt_opts`] but lets a test pin the disposition explicitly.
|
||||
fn prompt_opts_eager<'a>(
|
||||
writeback: bool,
|
||||
mode: Option<&'a str>,
|
||||
eagerness: Option<&'a str>,
|
||||
target: &'a str,
|
||||
) -> KnowledgeContextOptions<'a> {
|
||||
KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::PromptSection,
|
||||
writeback,
|
||||
writeback_mode: mode,
|
||||
writeback_eagerness: eagerness,
|
||||
target_id: target,
|
||||
has_search_tool: false,
|
||||
has_write_tool: false,
|
||||
}
|
||||
}
|
||||
|
||||
// ── empty input ──────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn empty_mounts_build_nothing() {
|
||||
assert_eq!(build_knowledge_context(&[], &prompt_opts(false, None, "conv-1")), None);
|
||||
let readme_opts = KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::TerminalReadme,
|
||||
writeback: true,
|
||||
writeback_mode: None,
|
||||
writeback_eagerness: None,
|
||||
target_id: "term-1",
|
||||
has_search_tool: false,
|
||||
has_write_tool: false,
|
||||
};
|
||||
assert_eq!(build_knowledge_context(&[], &readme_opts), None);
|
||||
}
|
||||
|
||||
// ── single base: full per-base contract ──────────────────────────
|
||||
|
||||
#[test]
|
||||
fn single_base_prompt_section_contract() {
|
||||
let mut m = mount("领域知识", ".nomi/knowledge/领域知识");
|
||||
m.description = "团队约定".into();
|
||||
m.summary = Some("Covers deployment flows and on-call runbooks.".into());
|
||||
m.toc = vec!["concepts/术语.md — 术语表".into(), "(+3 more files)".into()];
|
||||
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
|
||||
// Section heading stays compatible with the historical one.
|
||||
assert!(out.starts_with("## Knowledge bases (extended knowledge source)"), "got: {out}");
|
||||
// Retrieval protocol: search-before-answer, Grep/Read guidance,
|
||||
// relative-path citation, abridged-TOC explore hint — rendered once.
|
||||
assert!(out.contains("Retrieval protocol"), "got: {out}");
|
||||
assert!(out.contains("BEFORE answering"), "got: {out}");
|
||||
assert!(out.contains("Grep"), "got: {out}");
|
||||
assert!(out.contains("relative path"), "got: {out}");
|
||||
assert!(out.contains("summarizes a folder"), "got: {out}");
|
||||
// Per-base section: name, path, description, summary, when-to-consult, TOC.
|
||||
assert!(out.contains("### 领域知识"), "got: {out}");
|
||||
assert!(out.contains("`./.nomi/knowledge/领域知识/`"), "got: {out}");
|
||||
assert!(out.contains("团队约定"), "got: {out}");
|
||||
assert!(out.contains("Covers deployment flows and on-call runbooks."), "got: {out}");
|
||||
assert!(out.contains("When to consult"), "got: {out}");
|
||||
assert!(out.contains("concepts/术语.md — 术语表"), "got: {out}");
|
||||
assert!(out.contains("(+3 more files)"), "got: {out}");
|
||||
// Read-only contract when write-back is off.
|
||||
assert!(out.contains("Write-back is DISABLED"), "got: {out}");
|
||||
assert!(out.contains("READ-ONLY"), "got: {out}");
|
||||
// No live sources → no realtime section, no fetch-tool mention.
|
||||
assert!(!out.contains("Realtime sources"), "got: {out}");
|
||||
assert!(!out.contains("nomi_knowledge_fetch_url"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_description_and_summary_lines_are_omitted() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert!(!out.contains("Description:"), "got: {out}");
|
||||
assert!(!out.contains("Summary:"), "got: {out}");
|
||||
// The when-to-consult guidance survives even without description.
|
||||
assert!(out.contains("When to consult"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_description_base_surfaces_topic_hints_from_toc() {
|
||||
let mut m = mount("库A", ".nomi/knowledge/库A");
|
||||
m.toc = vec![
|
||||
"deploy/rollback.md — 回滚流程".into(),
|
||||
"concepts/术语.md — 术语表".into(),
|
||||
"docs/ — 12 files".into(),
|
||||
"(+3 more files)".into(),
|
||||
];
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert!(out.contains("Topics include:"), "got: {out}");
|
||||
assert!(out.contains("回滚流程"), "got: {out}");
|
||||
assert!(out.contains("术语表"), "got: {out}");
|
||||
// Aggregate rows (`dir/ — N files`) and the `(+N more)` remainder must
|
||||
// not be promoted to hints. The verbatim TOC still lists them under
|
||||
// `- Contents:` and the protocol preamble names one as an example, so
|
||||
// scope the check to the `Topics include:` line itself.
|
||||
let hint_line = out
|
||||
.lines()
|
||||
.find(|l| l.starts_with("- Topics include:"))
|
||||
.expect("hint line present");
|
||||
assert!(!hint_line.contains("12 files"), "aggregate rows must not be hints: {hint_line}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn described_base_omits_topic_hints() {
|
||||
let mut m = mount("库A", ".nomi/knowledge/库A");
|
||||
m.description = "团队约定".into();
|
||||
m.toc = vec!["x.md — 标题".into()];
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert!(!out.contains("Topics include:"), "described base needs no hint line: {out}");
|
||||
assert!(out.contains("团队约定"));
|
||||
}
|
||||
|
||||
// ── multi base: protocol rendered once ───────────────────────────
|
||||
|
||||
#[test]
|
||||
fn multi_base_renders_protocol_once() {
|
||||
let a = mount("库A", ".nomi/knowledge/库A");
|
||||
let b = mount("库B", ".nomi/knowledge/库B");
|
||||
let out = build_knowledge_context(&[a, b], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert_eq!(out.matches("Retrieval protocol").count(), 1, "got: {out}");
|
||||
assert_eq!(out.matches("Write-back is DISABLED").count(), 1, "got: {out}");
|
||||
assert!(out.contains("### 库A"), "got: {out}");
|
||||
assert!(out.contains("### 库B"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── tool-aware retrieval protocol ────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn protocol_mentions_search_tool_when_available() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let mut opts = prompt_opts(false, None, "conv-1");
|
||||
opts.has_search_tool = true;
|
||||
let out = build_knowledge_context(std::slice::from_ref(&m), &opts).unwrap();
|
||||
assert!(out.contains("call the `knowledge_search` tool"), "got: {out}");
|
||||
assert!(!out.contains("Grep/Glob instead of"), "search variant drops Grep-first wording: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_keeps_grep_wording_without_search_tool() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert!(out.contains("Grep/Glob"), "got: {out}");
|
||||
assert!(!out.contains("knowledge_search"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── TOC budgets ──────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn toc_budget_keeps_small_listings_untouched() {
|
||||
let mut tocs = vec![vec!["a.md".to_string(), "b.md — B".to_string()]];
|
||||
apply_toc_budgets(&mut tocs);
|
||||
assert_eq!(tocs[0], vec!["a.md".to_string(), "b.md — B".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toc_budget_aggregates_per_kb_overflow_by_directory() {
|
||||
// 10 root files + 15 under docs/ = 25 sorted lines; per-KB budget 20
|
||||
// → keep first 20, aggregate the 5 overflowing docs/ files.
|
||||
let mut lines: Vec<String> = (0..10).map(|i| format!("a{i:02}.md — Root {i}")).collect();
|
||||
lines.extend((0..15).map(|i| format!("docs/d{i:02}.md — Doc {i}")));
|
||||
let mut tocs = vec![lines];
|
||||
apply_toc_budgets(&mut tocs);
|
||||
|
||||
let toc = &tocs[0];
|
||||
assert_eq!(toc.len(), TOC_PER_KB_MAX + 1, "got: {toc:?}");
|
||||
assert_eq!(toc[0], "a00.md — Root 0");
|
||||
assert_eq!(toc[TOC_PER_KB_MAX - 1], "docs/d09.md — Doc 9");
|
||||
assert_eq!(toc[TOC_PER_KB_MAX], "docs/ — 5 files");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toc_budget_rootless_overflow_keeps_more_files_marker() {
|
||||
let mut tocs = vec![(0..25).map(|i| format!("f{i:02}.md")).collect::<Vec<_>>()];
|
||||
apply_toc_budgets(&mut tocs);
|
||||
let toc = &tocs[0];
|
||||
assert_eq!(toc.len(), TOC_PER_KB_MAX + 1, "got: {toc:?}");
|
||||
assert_eq!(toc[TOC_PER_KB_MAX], "(+5 more files)");
|
||||
}
|
||||
|
||||
/// Many distinct overflowing directories must not balloon the TOC: only
|
||||
/// the top-[`TOC_AGGREGATE_DIR_MAX`] directories get named rows, the
|
||||
/// rest (plus rootless overflow) folds into `(+N more files)`.
|
||||
#[test]
|
||||
fn toc_budget_caps_aggregate_directory_rows() {
|
||||
// Budget-filling root files first, then 10 dirs with growing file
|
||||
// counts (d0: 1 file … d9: 10 files) plus 3 rootless files.
|
||||
let mut lines: Vec<String> = (0..TOC_PER_KB_MAX).map(|i| format!("a{i:02}.md")).collect();
|
||||
for d in 0..10 {
|
||||
for f in 0..=d {
|
||||
lines.push(format!("d{d}/f{f:02}.md"));
|
||||
}
|
||||
}
|
||||
lines.extend((0..3).map(|i| format!("z{i}.md")));
|
||||
let mut tocs = vec![lines];
|
||||
apply_toc_budgets(&mut tocs);
|
||||
|
||||
let toc = &tocs[0];
|
||||
let dir_rows: Vec<&String> = toc.iter().filter(|l| l.contains("/ — ")).collect();
|
||||
assert_eq!(dir_rows.len(), TOC_AGGREGATE_DIR_MAX, "got: {toc:?}");
|
||||
// The two SMALLEST dirs (d0: 1, d1: 2) are folded, the rest named.
|
||||
assert!(!toc.iter().any(|l| l.starts_with("d0/")), "got: {toc:?}");
|
||||
assert!(!toc.iter().any(|l| l.starts_with("d1/")), "got: {toc:?}");
|
||||
assert!(toc.contains(&"d9/ — 10 files".to_string()), "got: {toc:?}");
|
||||
// Named rows stay in path order.
|
||||
assert_eq!(dir_rows[0], "d2/ — 3 files", "got: {toc:?}");
|
||||
// Remainder: d0(1) + d1(2) dirs + 3 rootless = 6.
|
||||
assert_eq!(toc.last().unwrap(), "(+6 more files)", "got: {toc:?}");
|
||||
// Total rows = budget + 8 dir rows + 1 remainder row.
|
||||
assert_eq!(toc.len(), TOC_PER_KB_MAX + TOC_AGGREGATE_DIR_MAX + 1, "got: {toc:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toc_budget_shrinks_per_kb_share_under_global_cap() {
|
||||
// 4 bases × 20 files = 80 > 60 → fair share 15 file lines each.
|
||||
let mut tocs: Vec<Vec<String>> = (0..4)
|
||||
.map(|kb| (0..20).map(|i| format!("kb{kb}/f{i:02}.md")).collect())
|
||||
.collect();
|
||||
apply_toc_budgets(&mut tocs);
|
||||
for toc in &tocs {
|
||||
let file_lines = toc.iter().filter(|l| l.contains(".md")).count();
|
||||
assert_eq!(file_lines, 15, "got: {toc:?}");
|
||||
assert!(toc.iter().any(|l| l.ends_with("— 5 files")), "got: {toc:?}");
|
||||
}
|
||||
}
|
||||
|
||||
// ── write-back contract ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn staged_writeback_scopes_inbox_to_target_id() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
// Default mode (None) is staged.
|
||||
let out = build_knowledge_context(std::slice::from_ref(&m), &prompt_opts(true, None, "term_42")).unwrap();
|
||||
assert!(out.contains("STAGED mode"), "got: {out}");
|
||||
assert!(out.contains("_inbox/term_42/"), "got: {out}");
|
||||
assert!(out.contains("READ-ONLY"), "got: {out}");
|
||||
// Explicit "staged" renders identically.
|
||||
let explicit = build_knowledge_context(&[m], &prompt_opts(true, Some("staged"), "term_42")).unwrap();
|
||||
assert_eq!(out, explicit);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_writeback_never_mentions_inbox() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(true, Some("direct"), "conv-1")).unwrap();
|
||||
assert!(out.contains("DIRECT mode"), "got: {out}");
|
||||
assert!(!out.contains("_inbox"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_writeback_contract_directs_handle_use_for_updates() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
// DIRECT + tools available: update via handle, read via knowledge_read.
|
||||
let mut direct = prompt_opts(true, Some("direct"), "conv-1");
|
||||
direct.has_write_tool = true;
|
||||
direct.has_search_tool = true;
|
||||
let out = build_knowledge_context(std::slice::from_ref(&m), &direct).unwrap();
|
||||
assert!(out.contains("knowledge_write"), "got: {out}");
|
||||
assert!(out.contains("handle"), "update path must reference the handle: {out}");
|
||||
assert!(out.contains("knowledge_read"), "got: {out}");
|
||||
// STAGED + tools: emphasize auto-placement + original untouched.
|
||||
let mut staged = prompt_opts(true, Some("staged"), "conv-1");
|
||||
staged.has_write_tool = true;
|
||||
staged.has_search_tool = true;
|
||||
let s = build_knowledge_context(&[m], &staged).unwrap();
|
||||
assert!(s.contains("handle") && s.contains("review inbox"), "got: {s}");
|
||||
assert!(s.contains("left untouched"), "got: {s}");
|
||||
}
|
||||
|
||||
// ── realtime (live URL) sources ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn live_sources_render_realtime_section() {
|
||||
let mut m = mount("接口库", ".nomi/knowledge/接口库");
|
||||
m.live_sources = vec![
|
||||
KnowledgeSourceEntry {
|
||||
url: "https://example.com/api-docs".into(),
|
||||
title: Some("API docs".into()),
|
||||
rendered: false,
|
||||
},
|
||||
KnowledgeSourceEntry {
|
||||
url: "https://example.com/changelog".into(),
|
||||
title: None,
|
||||
rendered: false,
|
||||
},
|
||||
];
|
||||
let plain = mount("普通库", ".nomi/knowledge/普通库");
|
||||
|
||||
let out = build_knowledge_context(&[m, plain], &prompt_opts(false, None, "conv-1")).unwrap();
|
||||
assert!(out.contains("Realtime sources"), "got: {out}");
|
||||
assert!(out.contains("API docs"), "got: {out}");
|
||||
assert!(out.contains("https://example.com/api-docs"), "got: {out}");
|
||||
assert!(out.contains("https://example.com/changelog"), "got: {out}");
|
||||
assert!(out.contains("接口库"), "got: {out}");
|
||||
// Layered tool guidance: own web-fetch tools first, the gateway tool
|
||||
// only as a conditional option, and an honest no-tools fallback.
|
||||
assert!(out.contains("web-fetch tool already available"), "got: {out}");
|
||||
assert!(out.contains("if a `nomi_knowledge_fetch_url` tool is available"), "got: {out}");
|
||||
assert!(out.contains("cannot read realtime sources"), "got: {out}");
|
||||
// The old unconditional promise must be gone — the gateway tool only
|
||||
// exists in desktopGateway sessions.
|
||||
assert!(!out.contains("call the `nomi_knowledge_fetch_url` tool instead"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── writeback mode parsing ───────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn writeback_mode_parses_known_values_and_falls_back_to_staged() {
|
||||
assert_eq!(WritebackMode::parse(None), WritebackMode::Staged);
|
||||
assert_eq!(WritebackMode::parse(Some("staged")), WritebackMode::Staged);
|
||||
assert_eq!(WritebackMode::parse(Some("direct")), WritebackMode::Direct);
|
||||
// Unknown (or wrong-case) values must never pick the permissive mode.
|
||||
assert_eq!(WritebackMode::parse(Some("DIRECT")), WritebackMode::Staged);
|
||||
assert_eq!(WritebackMode::parse(Some("yolo")), WritebackMode::Staged);
|
||||
assert_eq!(WritebackMode::parse(Some("")), WritebackMode::Staged);
|
||||
assert_eq!(WritebackMode::default(), WritebackMode::Staged);
|
||||
}
|
||||
|
||||
/// An unrecognized writeback_mode string renders the STAGED contract —
|
||||
/// never DIRECT (the permissive branch must be opt-in by exact value).
|
||||
#[test]
|
||||
fn unknown_writeback_mode_renders_staged_contract() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(&[m], &prompt_opts(true, Some("yolo"), "conv-7")).unwrap();
|
||||
assert!(out.contains("STAGED mode"), "got: {out}");
|
||||
assert!(out.contains("_inbox/conv-7/"), "got: {out}");
|
||||
assert!(!out.contains("DIRECT mode"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── writeback eagerness (回写意识) ─────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn writeback_eagerness_parses_known_values_and_falls_back_to_conservative() {
|
||||
assert_eq!(WritebackEagerness::parse(None), WritebackEagerness::Conservative);
|
||||
assert_eq!(WritebackEagerness::parse(Some("conservative")), WritebackEagerness::Conservative);
|
||||
assert_eq!(WritebackEagerness::parse(Some("aggressive")), WritebackEagerness::Aggressive);
|
||||
// Unknown / wrong-case values must never pick the eager mode.
|
||||
assert_eq!(WritebackEagerness::parse(Some("AGGRESSIVE")), WritebackEagerness::Conservative);
|
||||
assert_eq!(WritebackEagerness::parse(Some("bold")), WritebackEagerness::Conservative);
|
||||
assert_eq!(WritebackEagerness::parse(Some("")), WritebackEagerness::Conservative);
|
||||
assert_eq!(WritebackEagerness::default(), WritebackEagerness::Conservative);
|
||||
}
|
||||
|
||||
/// Enabled write-back defaults to the CONSERVATIVE disposition clause and
|
||||
/// is orthogonal to the staged/direct placement decision.
|
||||
#[test]
|
||||
fn enabled_writeback_appends_conservative_clause_by_default() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
// Default eagerness (None) under staged mode.
|
||||
let staged = build_knowledge_context(&[m.clone()], &prompt_opts(true, None, "conv-1")).unwrap();
|
||||
assert!(staged.contains("STAGED mode"), "got: {staged}");
|
||||
assert!(staged.contains("Disposition — CONSERVATIVE"), "got: {staged}");
|
||||
assert!(!staged.contains("Disposition — AGGRESSIVE"), "got: {staged}");
|
||||
// Default eagerness under direct mode too.
|
||||
let direct = build_knowledge_context(&[m], &prompt_opts(true, Some("direct"), "conv-1")).unwrap();
|
||||
assert!(direct.contains("DIRECT mode"), "got: {direct}");
|
||||
assert!(direct.contains("Disposition — CONSERVATIVE"), "got: {direct}");
|
||||
}
|
||||
|
||||
/// The aggressive disposition renders its own clause, independent of mode.
|
||||
#[test]
|
||||
fn aggressive_eagerness_renders_aggressive_clause_for_both_modes() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let staged = build_knowledge_context(
|
||||
&[m.clone()],
|
||||
&prompt_opts_eager(true, Some("staged"), Some("aggressive"), "conv-1"),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(staged.contains("STAGED mode"), "got: {staged}");
|
||||
assert!(staged.contains("Disposition — AGGRESSIVE"), "got: {staged}");
|
||||
assert!(!staged.contains("Disposition — CONSERVATIVE"), "got: {staged}");
|
||||
// Staged placement survives an aggressive disposition (inbox still scoped).
|
||||
assert!(staged.contains("_inbox/conv-1/"), "got: {staged}");
|
||||
|
||||
let direct = build_knowledge_context(
|
||||
&[m],
|
||||
&prompt_opts_eager(true, Some("direct"), Some("aggressive"), "conv-1"),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(direct.contains("DIRECT mode"), "got: {direct}");
|
||||
assert!(direct.contains("Disposition — AGGRESSIVE"), "got: {direct}");
|
||||
assert!(!direct.contains("_inbox"), "got: {direct}");
|
||||
}
|
||||
|
||||
/// Disabled write-back is read-only and carries no disposition clause —
|
||||
/// eagerness is meaningless without write-back.
|
||||
#[test]
|
||||
fn disabled_writeback_has_no_eagerness_clause() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(
|
||||
&[m],
|
||||
&prompt_opts_eager(false, None, Some("aggressive"), "conv-1"),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(out.contains("Write-back is DISABLED"), "got: {out}");
|
||||
assert!(!out.contains("Disposition —"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── tool-based write-back contract (has_write_tool = true) ────────
|
||||
|
||||
/// Options helper pinning has_write_tool = true (the nomi-engine surface).
|
||||
fn prompt_opts_tooled<'a>(mode: Option<&'a str>, target: &'a str) -> KnowledgeContextOptions<'a> {
|
||||
KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::PromptSection,
|
||||
writeback: true,
|
||||
writeback_mode: mode,
|
||||
writeback_eagerness: None,
|
||||
target_id: target,
|
||||
has_search_tool: true,
|
||||
has_write_tool: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// When the surface has the native tool, the contract instructs CALLING
|
||||
/// `knowledge_write` and drops the file-path / inbox-path prose — in BOTH
|
||||
/// modes. The model must never be pointed at the generic Write tool.
|
||||
#[test]
|
||||
fn tool_contract_directs_to_knowledge_write_in_both_modes() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let staged = build_knowledge_context(&[m.clone()], &prompt_opts_tooled(Some("staged"), "conv-1")).unwrap();
|
||||
assert!(staged.contains("knowledge_write"), "got: {staged}");
|
||||
assert!(staged.contains("STAGED mode"), "got: {staged}");
|
||||
assert!(staged.contains("Do NOT use the generic Write/Edit"), "got: {staged}");
|
||||
// The staged inbox PATH is now internal to the tool — never leaked to the model.
|
||||
assert!(!staged.contains("_inbox/"), "tool contract must not advertise the inbox path: {staged}");
|
||||
|
||||
let direct = build_knowledge_context(&[m], &prompt_opts_tooled(Some("direct"), "conv-1")).unwrap();
|
||||
assert!(direct.contains("knowledge_write"), "got: {direct}");
|
||||
assert!(direct.contains("DIRECT mode"), "got: {direct}");
|
||||
assert!(direct.contains("Do NOT use the generic Write/Edit"), "got: {direct}");
|
||||
// Disposition clause still appends under the tool contract.
|
||||
assert!(direct.contains("Disposition — CONSERVATIVE"), "got: {direct}");
|
||||
}
|
||||
|
||||
/// Disabled write-back is read-only regardless of has_write_tool.
|
||||
#[test]
|
||||
fn tool_surface_still_read_only_when_writeback_disabled() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let opts = KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::PromptSection,
|
||||
writeback: false,
|
||||
writeback_mode: None,
|
||||
writeback_eagerness: None,
|
||||
target_id: "conv-1",
|
||||
has_search_tool: true,
|
||||
has_write_tool: true,
|
||||
};
|
||||
let out = build_knowledge_context(&[m], &opts).unwrap();
|
||||
assert!(out.contains("Write-back is DISABLED"), "got: {out}");
|
||||
assert!(!out.contains("knowledge_write"), "got: {out}");
|
||||
}
|
||||
|
||||
/// An unrecognized eagerness string renders the CONSERVATIVE clause — the
|
||||
/// eager branch must be opt-in by exact value, never reached by a typo.
|
||||
#[test]
|
||||
fn unknown_eagerness_renders_conservative_clause() {
|
||||
let m = mount("库A", ".nomi/knowledge/库A");
|
||||
let out = build_knowledge_context(
|
||||
&[m],
|
||||
&prompt_opts_eager(true, None, Some("yolo"), "conv-1"),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(out.contains("Disposition — CONSERVATIVE"), "got: {out}");
|
||||
assert!(!out.contains("Disposition — AGGRESSIVE"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── TerminalReadme format ────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn terminal_readme_is_a_complete_document() {
|
||||
let mut m = mount("领域知识", ".nomi/knowledge/领域知识");
|
||||
m.toc = vec!["intro.md — 简介".into()];
|
||||
let opts = KnowledgeContextOptions {
|
||||
format: KnowledgeContextFormat::TerminalReadme,
|
||||
writeback: true,
|
||||
writeback_mode: None,
|
||||
writeback_eagerness: None,
|
||||
target_id: "conv-9",
|
||||
has_search_tool: false,
|
||||
has_write_tool: false,
|
||||
};
|
||||
let out = build_knowledge_context(&[m], &opts).unwrap();
|
||||
|
||||
assert!(out.starts_with("# Knowledge bases"), "got: {out}");
|
||||
assert!(out.contains("NomiFun"), "got: {out}");
|
||||
// Relative-path baseline disambiguation (paths are workspace-rooted).
|
||||
assert!(out.contains("Paths below are relative to the workspace root."), "got: {out}");
|
||||
assert!(out.contains("## Retrieval protocol"), "got: {out}");
|
||||
assert!(out.contains("## Mounted bases"), "got: {out}");
|
||||
assert!(out.contains("### 领域知识"), "got: {out}");
|
||||
assert!(out.contains("intro.md — 简介"), "got: {out}");
|
||||
assert!(out.contains("STAGED mode"), "got: {out}");
|
||||
assert!(out.contains("_inbox/conv-9/"), "got: {out}");
|
||||
// The prompt-section heading must not leak into the readme format.
|
||||
assert!(!out.contains("## Knowledge bases (extended knowledge source)"), "got: {out}");
|
||||
}
|
||||
|
||||
// ── serde backward compatibility of the extra shape ──────────────
|
||||
|
||||
#[test]
|
||||
fn mount_info_deserializes_legacy_extra_without_new_fields() {
|
||||
let legacy = serde_json::json!({
|
||||
"id": "kb1",
|
||||
"name": "运维手册",
|
||||
"description": "",
|
||||
"rel_path": ".nomi/knowledge/运维手册",
|
||||
"toc": ["deploy.md — 部署"],
|
||||
});
|
||||
let m: KnowledgeMountInfo = serde_json::from_value(legacy).expect("legacy extra must deserialize");
|
||||
assert_eq!(m.summary, None);
|
||||
assert!(m.live_sources.is_empty());
|
||||
|
||||
// New fields round-trip, and empty optionals stay off the wire.
|
||||
let bare = serde_json::to_value(mount("库A", ".nomi/knowledge/库A")).unwrap();
|
||||
assert!(bare.get("summary").is_none(), "got: {bare}");
|
||||
assert!(bare.get("live_sources").is_none(), "got: {bare}");
|
||||
|
||||
let mut rich = mount("库B", ".nomi/knowledge/库B");
|
||||
rich.summary = Some("s".into());
|
||||
rich.live_sources = vec![KnowledgeSourceEntry {
|
||||
url: "https://e.com".into(),
|
||||
title: None,
|
||||
rendered: false,
|
||||
}];
|
||||
let v = serde_json::to_value(&rich).unwrap();
|
||||
let back: KnowledgeMountInfo = serde_json::from_value(v).unwrap();
|
||||
assert_eq!(back.summary.as_deref(), Some("s"));
|
||||
assert_eq!(back.live_sources.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
//! WS push events for the knowledge domain. Same shape as `CompanionEventEmitter`:
|
||||
//! a thin wrapper over the global `EventBroadcaster`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use nomifun_api_types::WebSocketMessage;
|
||||
use nomifun_realtime::EventBroadcaster;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct KnowledgeEventEmitter {
|
||||
broadcaster: Arc<dyn EventBroadcaster>,
|
||||
}
|
||||
|
||||
impl KnowledgeEventEmitter {
|
||||
pub fn new(broadcaster: Arc<dyn EventBroadcaster>) -> Self {
|
||||
Self { broadcaster }
|
||||
}
|
||||
|
||||
fn broadcast<T: serde::Serialize>(&self, event_name: &str, payload: &T) {
|
||||
let value = match serde_json::to_value(payload) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, event_name, "failed to serialize knowledge event");
|
||||
return;
|
||||
}
|
||||
};
|
||||
self.broadcaster.broadcast(WebSocketMessage::new(event_name, value));
|
||||
}
|
||||
|
||||
pub fn emit_base_created<T: serde::Serialize>(&self, base: &T) {
|
||||
self.broadcast("knowledge.base-created", base);
|
||||
}
|
||||
|
||||
pub fn emit_base_updated<T: serde::Serialize>(&self, base: &T) {
|
||||
self.broadcast("knowledge.base-updated", base);
|
||||
}
|
||||
|
||||
pub fn emit_base_deleted(&self, id: &str) {
|
||||
self.broadcast("knowledge.base-deleted", &serde_json::json!({ "id": id }));
|
||||
}
|
||||
|
||||
pub fn emit_binding_changed<T: serde::Serialize>(&self, binding: &T) {
|
||||
self.broadcast("knowledge.binding-changed", binding);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,620 @@
|
||||
//! Knowledge base zip export/import (spec §4.8) — cross-machine migration.
|
||||
//!
|
||||
//! Package layout (zip root):
|
||||
//! - `manifest.json` — `{format, version, kind, exported_at, app_version}`
|
||||
//! envelope, validated on import (wrong format/kind or a newer version is
|
||||
//! rejected before anything touches the registry).
|
||||
//! - `meta.json` — `{name, description}` of the exported base.
|
||||
//! - `files/**` — every `.md` under the base root with relative paths
|
||||
//! preserved. `_inbox/` is included on purpose: staged write-backs are
|
||||
//! user data and must survive a machine migration.
|
||||
//!
|
||||
//! Import extracts into a temp dir under the managed knowledge dir (same
|
||||
//! volume as the final destination, so file moves are cheap renames), with
|
||||
//! zip-slip hardening mirroring `nomifun-extension`'s skill import: entry
|
||||
//! paths are component-sanitized, symlink entries are rejected, and only
|
||||
//! `manifest.json` / `meta.json` / `files/**.md` entries are accepted.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::io::Write;
|
||||
use std::path::{Component, Path, PathBuf};
|
||||
|
||||
use nomifun_common::{AppError, TimestampMs, now_ms};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::KB_MANAGED_REL_DIR;
|
||||
use crate::service::{KnowledgeService, is_md};
|
||||
|
||||
/// `manifest.json` envelope discriminators. `version` is bumped only on
|
||||
/// breaking package-layout changes; readers accept anything `<= EXPORT_VERSION`.
|
||||
pub const EXPORT_FORMAT: &str = "nomifun-export";
|
||||
pub const EXPORT_KIND: &str = "knowledge-base";
|
||||
pub const EXPORT_VERSION: u32 = 1;
|
||||
|
||||
/// Result of a successful export, returned to the frontend.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ExportSummary {
|
||||
pub file_count: u64,
|
||||
/// Uncompressed size of the packaged `.md` files.
|
||||
pub total_bytes: u64,
|
||||
pub dest_path: String,
|
||||
}
|
||||
|
||||
/// Result of a successful import, returned to the frontend.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ImportSummary {
|
||||
pub kb_id: String,
|
||||
/// Final name after duplicate-name suffixing (`"name (2)"`, …).
|
||||
pub name: String,
|
||||
pub file_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ExportManifest {
|
||||
format: String,
|
||||
version: u32,
|
||||
kind: String,
|
||||
exported_at: TimestampMs,
|
||||
app_version: String,
|
||||
}
|
||||
|
||||
/// `meta.json` payload. Lenient on read: missing fields default to empty so
|
||||
/// a hand-edited package still imports.
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
struct ExportMeta {
|
||||
#[serde(default)]
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
}
|
||||
|
||||
// ── Export ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Package the base `kb_id` into a zip at `dest_path` (written atomically
|
||||
/// via `{dest}.tmp` + rename).
|
||||
pub async fn export_base(
|
||||
service: &KnowledgeService,
|
||||
kb_id: &str,
|
||||
dest_path: &Path,
|
||||
) -> Result<ExportSummary, AppError> {
|
||||
let info = service.get_base_info(kb_id).await?;
|
||||
if !info.root_exists {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"knowledge base directory missing: {}",
|
||||
info.root_path
|
||||
)));
|
||||
}
|
||||
if !dest_path.is_absolute() {
|
||||
return Err(AppError::BadRequest("dest_path must be absolute".into()));
|
||||
}
|
||||
|
||||
let root = PathBuf::from(&info.root_path);
|
||||
let meta = ExportMeta {
|
||||
name: info.name,
|
||||
description: info.description,
|
||||
};
|
||||
let dest = dest_path.to_path_buf();
|
||||
let (file_count, total_bytes) = tokio::task::spawn_blocking(move || build_zip(&root, &meta, &dest))
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(format!("export task join error: {e}")))??;
|
||||
|
||||
Ok(ExportSummary {
|
||||
file_count,
|
||||
total_bytes,
|
||||
dest_path: dest_path.to_string_lossy().to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Blocking core of the export: walk `root` for `.md` files and write the
|
||||
/// package to `dest` via a `.tmp` sibling. Returns `(file_count, total_bytes)`.
|
||||
fn build_zip(root: &Path, meta: &ExportMeta, dest: &Path) -> Result<(u64, u64), AppError> {
|
||||
if let Some(parent) = dest.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| AppError::Internal(format!("failed to create export dir: {e}")))?;
|
||||
}
|
||||
let mut tmp_name = dest.as_os_str().to_owned();
|
||||
tmp_name.push(".tmp");
|
||||
let tmp = PathBuf::from(tmp_name);
|
||||
|
||||
let counts = match write_zip_to(root, meta, &tmp) {
|
||||
Ok(counts) => counts,
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
if let Err(e) = std::fs::rename(&tmp, dest) {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(AppError::Internal(format!("failed to finalize export file: {e}")));
|
||||
}
|
||||
Ok(counts)
|
||||
}
|
||||
|
||||
fn write_zip_to(root: &Path, meta: &ExportMeta, tmp: &Path) -> Result<(u64, u64), AppError> {
|
||||
let io_err = |what: &str| {
|
||||
let what = what.to_owned();
|
||||
move |e: std::io::Error| AppError::Internal(format!("{what}: {e}"))
|
||||
};
|
||||
let zip_err = |e: zip::result::ZipError| AppError::Internal(format!("failed to write zip: {e}"));
|
||||
|
||||
let file = std::fs::File::create(tmp).map_err(io_err("failed to create export file"))?;
|
||||
let mut zip = zip::ZipWriter::new(file);
|
||||
let options = zip::write::SimpleFileOptions::default();
|
||||
|
||||
let manifest = ExportManifest {
|
||||
format: EXPORT_FORMAT.to_owned(),
|
||||
version: EXPORT_VERSION,
|
||||
kind: EXPORT_KIND.to_owned(),
|
||||
exported_at: now_ms(),
|
||||
app_version: env!("CARGO_PKG_VERSION").to_owned(),
|
||||
};
|
||||
zip.start_file("manifest.json", options).map_err(zip_err)?;
|
||||
zip.write_all(&serde_json::to_vec_pretty(&manifest).map_err(|e| AppError::Internal(e.to_string()))?)
|
||||
.map_err(io_err("failed to write manifest"))?;
|
||||
zip.start_file("meta.json", options).map_err(zip_err)?;
|
||||
zip.write_all(&serde_json::to_vec_pretty(meta).map_err(|e| AppError::Internal(e.to_string()))?)
|
||||
.map_err(io_err("failed to write meta"))?;
|
||||
|
||||
// Sorted relative paths → deterministic packages (friendlier diffing).
|
||||
let mut rels: Vec<String> = walkdir::WalkDir::new(root)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|e| e.file_type().is_file() && is_md(e.path()))
|
||||
.filter_map(|e| {
|
||||
e.path()
|
||||
.strip_prefix(root)
|
||||
.ok()
|
||||
.map(|rel| rel.to_string_lossy().replace('\\', "/"))
|
||||
})
|
||||
.collect();
|
||||
rels.sort();
|
||||
|
||||
let mut file_count = 0u64;
|
||||
let mut total_bytes = 0u64;
|
||||
for rel in rels {
|
||||
let bytes = std::fs::read(root.join(&rel)).map_err(io_err(&format!("failed to read {rel}")))?;
|
||||
zip.start_file(format!("files/{rel}"), options).map_err(zip_err)?;
|
||||
zip.write_all(&bytes).map_err(io_err(&format!("failed to package {rel}")))?;
|
||||
file_count += 1;
|
||||
total_bytes += bytes.len() as u64;
|
||||
}
|
||||
|
||||
zip.finish().map_err(zip_err)?;
|
||||
Ok((file_count, total_bytes))
|
||||
}
|
||||
|
||||
// ── Import ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Import a package created by [`export_base`]: validate, create a new
|
||||
/// managed base (name deduplicated against existing bases), and move the
|
||||
/// packaged files into its root. Emits `knowledge.base-created` via the
|
||||
/// service's create path, then `knowledge.base-updated` once files landed
|
||||
/// so clients see correct stats.
|
||||
pub async fn import_base(service: &KnowledgeService, src_path: &Path) -> Result<ImportSummary, AppError> {
|
||||
if !src_path.is_file() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"import file does not exist: {}",
|
||||
src_path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
// Extraction temp lives next to the managed bases (same volume → the
|
||||
// final move is a cheap rename), namespaced to avoid collisions.
|
||||
let tmp_root = service.data_dir().join(KB_MANAGED_REL_DIR).join(".import-tmp");
|
||||
let extract_dir = tmp_root.join(format!("kb-{}-{}", std::process::id(), now_ms()));
|
||||
tokio::fs::create_dir_all(&extract_dir)
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(format!("failed to create import temp dir: {e}")))?;
|
||||
|
||||
let result = import_extracted(service, src_path, &extract_dir).await;
|
||||
let _ = tokio::fs::remove_dir_all(&extract_dir).await;
|
||||
let _ = tokio::fs::remove_dir(&tmp_root).await; // best-effort, only when empty
|
||||
result
|
||||
}
|
||||
|
||||
async fn import_extracted(
|
||||
service: &KnowledgeService,
|
||||
src_path: &Path,
|
||||
extract_dir: &Path,
|
||||
) -> Result<ImportSummary, AppError> {
|
||||
let src = src_path.to_path_buf();
|
||||
let dest = extract_dir.to_path_buf();
|
||||
let meta = tokio::task::spawn_blocking(move || extract_zip_validated(&src, &dest))
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(format!("import task join error: {e}")))??;
|
||||
|
||||
let existing: HashSet<String> = service
|
||||
.list_bases()
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|info| info.name)
|
||||
.collect();
|
||||
let base_name = match meta.name.trim() {
|
||||
"" => "导入的知识库",
|
||||
name => name,
|
||||
};
|
||||
let final_name = dedup_name(&existing, base_name);
|
||||
|
||||
// Existing managed-create path: provisions `{data_dir}/knowledge/{id}`
|
||||
// and emits `knowledge.base-created`. (Imported packages carry no URL
|
||||
// source — `extra` starts empty.)
|
||||
let info = service.create_base(&final_name, &meta.description, None, None).await?;
|
||||
|
||||
let files_src = extract_dir.join("files");
|
||||
let files_dest = PathBuf::from(&info.root_path);
|
||||
let moved = tokio::task::spawn_blocking(move || move_file_tree(&files_src, &files_dest))
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(format!("import move task join error: {e}")));
|
||||
let file_count = match moved {
|
||||
Ok(Ok(count)) => count,
|
||||
Ok(Err(e)) | Err(e) => {
|
||||
// Roll back the half-created base (purge is safe: managed dir).
|
||||
if let Err(del) = service.delete_base(&info.id, true).await {
|
||||
tracing::warn!(kb_id = %info.id, error = %del, "rollback of failed import left a stale base");
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// Re-emit with fresh file stats so clients don't show a 0-file base.
|
||||
if let Err(e) = service.update_base(&info.id, None, None, None).await {
|
||||
tracing::warn!(kb_id = %info.id, error = %e, "failed to refresh base info after import");
|
||||
}
|
||||
|
||||
Ok(ImportSummary {
|
||||
kb_id: info.id,
|
||||
name: final_name,
|
||||
file_count,
|
||||
})
|
||||
}
|
||||
|
||||
/// Blocking extraction with validation. Only `manifest.json`, `meta.json`
|
||||
/// and `files/**.md` entries are accepted; every entry path is sanitized
|
||||
/// (zip-slip) and symlink entries are rejected. Returns the parsed meta
|
||||
/// after the manifest passed format/kind/version checks.
|
||||
fn extract_zip_validated(archive_path: &Path, destination: &Path) -> Result<ExportMeta, AppError> {
|
||||
let file = std::fs::File::open(archive_path)
|
||||
.map_err(|e| AppError::BadRequest(format!("failed to open import file: {e}")))?;
|
||||
let mut archive =
|
||||
zip::ZipArchive::new(file).map_err(|_| AppError::BadRequest("不是知识库导出包".into()))?;
|
||||
|
||||
for index in 0..archive.len() {
|
||||
let mut entry = archive
|
||||
.by_index(index)
|
||||
.map_err(|e| AppError::BadRequest(format!("corrupt zip archive: {e}")))?;
|
||||
let entry_name = entry.name().to_string();
|
||||
reject_zip_symlink(&entry, &entry_name)?;
|
||||
let rel = safe_zip_entry_path(&entry_name)?;
|
||||
|
||||
if entry.is_dir() {
|
||||
if !rel.starts_with("files") {
|
||||
return Err(AppError::BadRequest("不是知识库导出包".into()));
|
||||
}
|
||||
std::fs::create_dir_all(destination.join(&rel))
|
||||
.map_err(|e| AppError::Internal(format!("failed to extract dir: {e}")))?;
|
||||
continue;
|
||||
}
|
||||
|
||||
let allowed = rel == Path::new("manifest.json")
|
||||
|| rel == Path::new("meta.json")
|
||||
|| (rel.starts_with("files") && is_md(&rel));
|
||||
if !allowed {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"不是知识库导出包(包含不支持的条目: {entry_name})"
|
||||
)));
|
||||
}
|
||||
|
||||
let output_path = destination.join(&rel);
|
||||
// Defense in depth on top of component sanitization: the resolved
|
||||
// path must stay inside the extraction dir.
|
||||
if !output_path.starts_with(destination) {
|
||||
return Err(AppError::BadRequest(format!("非法压缩包条目: {entry_name}")));
|
||||
}
|
||||
if let Some(parent) = output_path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| AppError::Internal(format!("failed to extract dirs: {e}")))?;
|
||||
}
|
||||
let mut output = std::fs::File::create(&output_path)
|
||||
.map_err(|e| AppError::Internal(format!("failed to extract file: {e}")))?;
|
||||
std::io::copy(&mut entry, &mut output)
|
||||
.map_err(|e| AppError::Internal(format!("failed to extract file: {e}")))?;
|
||||
}
|
||||
|
||||
let manifest_bytes = std::fs::read(destination.join("manifest.json"))
|
||||
.map_err(|_| AppError::BadRequest("不是知识库导出包".into()))?;
|
||||
let manifest: serde_json::Value = serde_json::from_slice(&manifest_bytes)
|
||||
.map_err(|_| AppError::BadRequest("不是知识库导出包".into()))?;
|
||||
validate_manifest(&manifest)?;
|
||||
|
||||
let meta: ExportMeta = std::fs::read(destination.join("meta.json"))
|
||||
.ok()
|
||||
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
|
||||
.unwrap_or_default();
|
||||
Ok(meta)
|
||||
}
|
||||
|
||||
/// Envelope check. Parsed as loose JSON so future manifests with extra
|
||||
/// fields still pass — only `format`/`kind`/`version` are load-bearing.
|
||||
fn validate_manifest(manifest: &serde_json::Value) -> Result<(), AppError> {
|
||||
let format = manifest.get("format").and_then(|v| v.as_str());
|
||||
let kind = manifest.get("kind").and_then(|v| v.as_str());
|
||||
if format != Some(EXPORT_FORMAT) || kind != Some(EXPORT_KIND) {
|
||||
return Err(AppError::BadRequest("不是知识库导出包".into()));
|
||||
}
|
||||
let version = manifest.get("version").and_then(|v| v.as_u64()).unwrap_or(0);
|
||||
if version > u64::from(EXPORT_VERSION) {
|
||||
return Err(AppError::BadRequest("导入包版本过新,请升级应用".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Sanitize a zip entry name into a safe relative path (same policy as
|
||||
/// `nomifun-extension`'s skill import): no backslashes, no absolute paths,
|
||||
/// no `..`/prefix components.
|
||||
fn safe_zip_entry_path(name: &str) -> Result<PathBuf, AppError> {
|
||||
let invalid = || AppError::BadRequest(format!("非法压缩包条目: {name}"));
|
||||
if name.is_empty() || name.contains('\\') {
|
||||
return Err(invalid());
|
||||
}
|
||||
let path = Path::new(name);
|
||||
if path.is_absolute() {
|
||||
return Err(invalid());
|
||||
}
|
||||
let mut safe_path = PathBuf::new();
|
||||
for component in path.components() {
|
||||
match component {
|
||||
Component::Normal(part) => safe_path.push(part),
|
||||
Component::CurDir => {}
|
||||
_ => return Err(invalid()),
|
||||
}
|
||||
}
|
||||
if safe_path.as_os_str().is_empty() {
|
||||
return Err(invalid());
|
||||
}
|
||||
Ok(safe_path)
|
||||
}
|
||||
|
||||
fn reject_zip_symlink(entry: &zip::read::ZipFile<'_>, name: &str) -> Result<(), AppError> {
|
||||
if let Some(mode) = entry.unix_mode()
|
||||
&& mode & 0o170000 == 0o120000
|
||||
{
|
||||
return Err(AppError::BadRequest(format!("非法压缩包条目: {name}")));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Move every file under `src_root` to the same relative path under
|
||||
/// `dest_root` (rename with copy fallback). Missing `src_root` (a package
|
||||
/// with zero files) is fine. Returns the number of files moved.
|
||||
fn move_file_tree(src_root: &Path, dest_root: &Path) -> Result<u64, AppError> {
|
||||
if !src_root.is_dir() {
|
||||
return Ok(0);
|
||||
}
|
||||
let mut count = 0u64;
|
||||
for entry in walkdir::WalkDir::new(src_root) {
|
||||
let entry = entry.map_err(|e| AppError::Internal(format!("failed to walk imported files: {e}")))?;
|
||||
if !entry.file_type().is_file() {
|
||||
continue;
|
||||
}
|
||||
let rel = entry
|
||||
.path()
|
||||
.strip_prefix(src_root)
|
||||
.map_err(|e| AppError::Internal(format!("failed to relativize imported file: {e}")))?;
|
||||
let dest = dest_root.join(rel);
|
||||
if let Some(parent) = dest.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.map_err(|e| AppError::Internal(format!("failed to create import dirs: {e}")))?;
|
||||
}
|
||||
if std::fs::rename(entry.path(), &dest).is_err() {
|
||||
std::fs::copy(entry.path(), &dest)
|
||||
.map_err(|e| AppError::Internal(format!("failed to place imported file: {e}")))?;
|
||||
}
|
||||
count += 1;
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Suffix `name` with `" (2)"`, `" (3)"`, … until it no longer collides
|
||||
/// with an existing base name.
|
||||
fn dedup_name(existing: &HashSet<String>, name: &str) -> String {
|
||||
if !existing.contains(name) {
|
||||
return name.to_owned();
|
||||
}
|
||||
for n in 2u32.. {
|
||||
let candidate = format!("{name} ({n})");
|
||||
if !existing.contains(&candidate) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
unreachable!("u32 suffix space exhausted")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::testutil::make_service;
|
||||
|
||||
fn write_test_zip(path: &Path, entries: &[(&str, &str)]) {
|
||||
let file = std::fs::File::create(path).unwrap();
|
||||
let mut zip = zip::ZipWriter::new(file);
|
||||
let options = zip::write::SimpleFileOptions::default();
|
||||
for (name, content) in entries {
|
||||
zip.start_file(*name, options).unwrap();
|
||||
zip.write_all(content.as_bytes()).unwrap();
|
||||
}
|
||||
zip.finish().unwrap();
|
||||
}
|
||||
|
||||
fn manifest_json(version: u32, kind: &str) -> String {
|
||||
format!(
|
||||
r#"{{"format":"nomifun-export","version":{version},"kind":"{kind}","exported_at":0,"app_version":"0.0.0"}}"#
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn export_import_roundtrip_preserves_file_tree() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let source = make_service(&dir.path().join("data-a"));
|
||||
let kb = source.create_base("迁移源库", "换机测试", None, None).await.unwrap();
|
||||
source.write_file(&kb.id, "guide.md", "# 指南\n正文").await.unwrap();
|
||||
source.write_file(&kb.id, "sub/notes.md", "嵌套内容").await.unwrap();
|
||||
source
|
||||
.write_file(&kb.id, "_inbox/conv_x/draft.md", "# 草稿")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let zip_path = dir.path().join("out").join("kb.zip");
|
||||
let summary = export_base(&source, &kb.id, &zip_path).await.unwrap();
|
||||
assert_eq!(summary.file_count, 3);
|
||||
assert!(summary.total_bytes > 0);
|
||||
assert!(zip_path.is_file());
|
||||
assert!(!dir.path().join("out").join("kb.zip.tmp").exists(), "tmp must be renamed away");
|
||||
|
||||
// Import into a fresh service (the "new machine").
|
||||
let target = make_service(&dir.path().join("data-b"));
|
||||
let imported = import_base(&target, &zip_path).await.unwrap();
|
||||
assert_eq!(imported.name, "迁移源库");
|
||||
assert_eq!(imported.file_count, 3);
|
||||
|
||||
let original: Vec<String> = source
|
||||
.list_files(&kb.id)
|
||||
.await
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|f| f.rel_path)
|
||||
.collect();
|
||||
let restored: Vec<String> = target
|
||||
.list_files(&imported.kb_id)
|
||||
.await
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|f| f.rel_path)
|
||||
.collect();
|
||||
assert_eq!(original, restored);
|
||||
|
||||
let content = target.read_file(&imported.kb_id, "guide.md").await.unwrap();
|
||||
assert_eq!(content.content, "# 指南\n正文");
|
||||
let info = target.get_base_info(&imported.kb_id).await.unwrap();
|
||||
assert_eq!(info.description, "换机测试");
|
||||
assert!(info.managed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_rejects_zip_slip_entries() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let service = make_service(&dir.path().join("data"));
|
||||
let zip_path = dir.path().join("evil.zip");
|
||||
write_test_zip(
|
||||
&zip_path,
|
||||
&[
|
||||
("manifest.json", &manifest_json(1, EXPORT_KIND)),
|
||||
("meta.json", r#"{"name":"x","description":""}"#),
|
||||
("../evil.md", "escaped"),
|
||||
],
|
||||
);
|
||||
|
||||
let err = import_base(&service, &zip_path).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{err:?}");
|
||||
assert!(!dir.path().join("evil.md").exists());
|
||||
assert!(
|
||||
service.list_bases().await.unwrap().is_empty(),
|
||||
"no base may be created from a rejected package"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_rejects_wrong_kind_and_newer_version() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let service = make_service(&dir.path().join("data"));
|
||||
|
||||
let wrong_kind = dir.path().join("skills.zip");
|
||||
write_test_zip(
|
||||
&wrong_kind,
|
||||
&[
|
||||
("manifest.json", &manifest_json(1, "skill-pack")),
|
||||
("meta.json", r#"{"name":"x"}"#),
|
||||
],
|
||||
);
|
||||
let err = import_base(&service, &wrong_kind).await.unwrap_err();
|
||||
assert!(err.to_string().contains("不是知识库导出包"), "{err}");
|
||||
|
||||
let too_new = dir.path().join("future.zip");
|
||||
write_test_zip(
|
||||
&too_new,
|
||||
&[
|
||||
("manifest.json", &manifest_json(2, EXPORT_KIND)),
|
||||
("meta.json", r#"{"name":"x"}"#),
|
||||
],
|
||||
);
|
||||
let err = import_base(&service, &too_new).await.unwrap_err();
|
||||
assert!(err.to_string().contains("导入包版本过新"), "{err}");
|
||||
|
||||
let not_zip = dir.path().join("garbage.zip");
|
||||
std::fs::write(¬_zip, "definitely not a zip").unwrap();
|
||||
let err = import_base(&service, ¬_zip).await.unwrap_err();
|
||||
assert!(err.to_string().contains("不是知识库导出包"), "{err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_rejects_non_md_payload() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let service = make_service(&dir.path().join("data"));
|
||||
let zip_path = dir.path().join("exe.zip");
|
||||
write_test_zip(
|
||||
&zip_path,
|
||||
&[
|
||||
("manifest.json", &manifest_json(1, EXPORT_KIND)),
|
||||
("meta.json", r#"{"name":"x"}"#),
|
||||
("files/payload.exe", "MZ"),
|
||||
],
|
||||
);
|
||||
let err = import_base(&service, &zip_path).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{err:?}");
|
||||
assert!(service.list_bases().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_suffixes_duplicate_names() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let service = make_service(&dir.path().join("data"));
|
||||
service.create_base("我的库", "", None, None).await.unwrap();
|
||||
|
||||
let zip_path = dir.path().join("dup.zip");
|
||||
write_test_zip(
|
||||
&zip_path,
|
||||
&[
|
||||
("manifest.json", &manifest_json(1, EXPORT_KIND)),
|
||||
("meta.json", r#"{"name":"我的库","description":""}"#),
|
||||
("files/a.md", "# A"),
|
||||
],
|
||||
);
|
||||
|
||||
let first = import_base(&service, &zip_path).await.unwrap();
|
||||
assert_eq!(first.name, "我的库 (2)");
|
||||
let second = import_base(&service, &zip_path).await.unwrap();
|
||||
assert_eq!(second.name, "我的库 (3)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedup_name_picks_first_free_suffix() {
|
||||
let mut existing = HashSet::new();
|
||||
assert_eq!(dedup_name(&existing, "库"), "库");
|
||||
existing.insert("库".to_owned());
|
||||
assert_eq!(dedup_name(&existing, "库"), "库 (2)");
|
||||
existing.insert("库 (2)".to_owned());
|
||||
existing.insert("库 (3)".to_owned());
|
||||
assert_eq!(dedup_name(&existing, "库"), "库 (4)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn safe_zip_entry_path_policy() {
|
||||
assert!(safe_zip_entry_path("files/a.md").is_ok());
|
||||
assert!(safe_zip_entry_path("./files/a.md").is_ok());
|
||||
assert!(safe_zip_entry_path("../evil.md").is_err());
|
||||
assert!(safe_zip_entry_path("files/../../evil.md").is_err());
|
||||
assert!(safe_zip_entry_path("/abs.md").is_err());
|
||||
assert!(safe_zip_entry_path("files\\win.md").is_err());
|
||||
assert!(safe_zip_entry_path("").is_err());
|
||||
assert!(safe_zip_entry_path("C:/evil.md").is_err());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,65 @@
|
||||
//! `nomifun-knowledge` — the Knowledge Base platform domain: user-curated
|
||||
//! directories of markdown documents, registered globally and mounted
|
||||
//! (junction/symlink) into session workspaces as an extended knowledge
|
||||
//! source. Sessions with the write-back ("回血") switch enabled are told via
|
||||
//! prompt contract that they may persist new knowledge back into the mounted
|
||||
//! directories.
|
||||
//!
|
||||
//! Layering: `service` owns registry CRUD + file access + mount planning;
|
||||
//! `mount` is the platform-aware link engine (junction on Windows, symlink on
|
||||
//! Unix, recursive copy fallback); `routes`/`state` are the `/api/knowledge/*`
|
||||
//! surface; `events` pushes WS notifications.
|
||||
//!
|
||||
//! The directory is the source of truth for content — the database only
|
||||
//! stores registration metadata, so users may drop `.md` files in at any
|
||||
//! time. Consumers other than conversations (terminal, companion) reuse the same
|
||||
//! `(target_kind, target_id)` binding storage; the companion integration is
|
||||
//! intentionally deferred (no code here depends on conversation or companion).
|
||||
|
||||
pub mod autogen;
|
||||
pub mod connector;
|
||||
pub mod connector_feishu;
|
||||
pub mod context;
|
||||
pub mod events;
|
||||
pub mod export;
|
||||
pub mod feishu_md;
|
||||
pub mod mcp_server;
|
||||
pub mod mount;
|
||||
pub mod routes;
|
||||
pub mod service;
|
||||
pub mod source_url;
|
||||
pub mod state;
|
||||
pub mod workpath;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod testutil;
|
||||
|
||||
pub use autogen::KnowledgeCompleter;
|
||||
pub use context::{KnowledgeContextFormat, KnowledgeContextOptions, WritebackEagerness, WritebackMode, build_knowledge_context};
|
||||
pub use events::KnowledgeEventEmitter;
|
||||
pub use mcp_server::KnowledgeMcpServer;
|
||||
pub use routes::knowledge_routes;
|
||||
pub use service::{
|
||||
AutogenOutcome, ConsumerInfo, InboxDiff, InboxEntry, InboxMergeResult, KB_INBOX_REL_DIR, KnowledgeBinding,
|
||||
KnowledgeService, MountOutcome, RefreshSourceSummary, WriteMode, WriteOp, WriteOutcome, WritePolicy, WriteRequest,
|
||||
WriteResolution, WriteSurface, WriteTargetSpec, decode_doc_handle, encode_doc_handle, resolve_write_policy,
|
||||
};
|
||||
pub use source_url::{HttpFetcher, PageFetcher, UrlFetcher};
|
||||
pub use state::KnowledgeRouterState;
|
||||
pub use workpath::{DEFAULT_WORKPATH_KEY, WORKPATH_BINDING_KIND, session_workpath_key, workpath_key};
|
||||
|
||||
/// Workspace-relative directory where knowledge bases are mounted. Lives
|
||||
/// under the hidden `.nomi/` folder — the same agent-facing namespace as
|
||||
/// `.nomi/skills` / `.nomi/plans` — so mounting into a user's own project
|
||||
/// directory stays unobtrusive (the mount dir self-ignores via its own
|
||||
/// `.gitignore`, see `mount.rs`).
|
||||
pub const KB_MOUNT_REL_DIR: &str = ".nomi/knowledge";
|
||||
|
||||
/// Pre-`.nomi` mount location. Kept solely so `mount::sync_mounts` can sweep
|
||||
/// leftover links/scaffolding out of workspaces created before the rename —
|
||||
/// never mount anything here.
|
||||
pub const KB_LEGACY_MOUNT_REL_DIR: &str = ".nomifun/knowledge";
|
||||
|
||||
/// Subdirectory of the backend data dir that hosts managed base directories:
|
||||
/// `{data_dir}/knowledge/{kb_id}/`.
|
||||
pub const KB_MANAGED_REL_DIR: &str = "knowledge";
|
||||
@@ -0,0 +1,670 @@
|
||||
//! In-process HTTP MCP server exposing the single `knowledge_search` tool to
|
||||
//! ACP agent sessions (claude / codex / gemini CLIs).
|
||||
//!
|
||||
//! ## Why this exists
|
||||
//!
|
||||
//! AutoWork drives ACP sessions, but ACP CLIs have no in-process tool bus we
|
||||
//! can register the native `KnowledgeSearchTool` into (only the nomi engine
|
||||
//! does). To give ACP agents the same knowledge-retrieval surface the nomi
|
||||
//! engine has natively, this server exposes ONE scoped tool, `knowledge_search`,
|
||||
//! over authenticated HTTP. Scope resolution has two paths:
|
||||
//!
|
||||
//! 1. **Explicit `kb_ids`** — baked at injection time and forwarded by the stdio
|
||||
//! bridge in each request body. The model searches only those bases.
|
||||
//! 2. **Runtime `cwd` resolution** — when no explicit `kb_ids` are supplied, the
|
||||
//! server resolves scope from the caller's working directory: workpath-bound
|
||||
//! bases if an enabled binding exists, or all mounted bases as fallback.
|
||||
//!
|
||||
//! In both paths the security invariant holds: the model supplies only `query`;
|
||||
//! scope is decided server-side; the model cannot widen the searchable set.
|
||||
//!
|
||||
//! ## Shape (mirrors `nomifun-requirement::mcp_server::RequirementMcpServer`)
|
||||
//!
|
||||
//! This is the in-process HTTP half. ACP CLIs spawn a SEPARATE stdio process
|
||||
//! (`nomicore mcp-knowledge-stdio`) that cannot share this process's
|
||||
//! `KnowledgeService`; it forwards each tool call back here as an authenticated
|
||||
//! `POST /tool`. The transport is stdio because claude / codex / gemini
|
||||
//! advertise stdio-only MCP capabilities (HTTP/SSE servers are dropped by the
|
||||
//! ACP capability filter), so a direct-HTTP injection would never reach them.
|
||||
//!
|
||||
//! ## Security
|
||||
//!
|
||||
//! A random opaque bearer token gates every request (per-process, like the
|
||||
//! requirement server). The tool is read-only, so there is no mutation scope to
|
||||
//! verify beyond the bound base set carried in `kb_ids`.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use axum::Json;
|
||||
use axum::extract::State;
|
||||
use axum::http::{HeaderValue, StatusCode, header};
|
||||
use axum::response::IntoResponse;
|
||||
use nomifun_common::generate_id;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::service::{
|
||||
KnowledgeBinding, KnowledgeSearchHit, KnowledgeService, WriteOp, WriteRequest, WriteSurface, WriteTargetSpec,
|
||||
decode_doc_handle, encode_doc_handle, resolve_write_policy,
|
||||
};
|
||||
|
||||
/// Late-bound handle to the singleton `KnowledgeService`. Held as a `Weak` so
|
||||
/// the server never keeps the service alive on its own (matches the requirement
|
||||
/// server's slot pattern). Wired via [`KnowledgeMcpServer::set_service`].
|
||||
type ServiceSlot = Arc<RwLock<Weak<KnowledgeService>>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct KbMcpState {
|
||||
auth_token: String,
|
||||
service: ServiceSlot,
|
||||
}
|
||||
|
||||
/// In-process HTTP MCP server for the scoped `knowledge_search` tool.
|
||||
pub struct KnowledgeMcpServer {
|
||||
http_addr: SocketAddr,
|
||||
auth_token: String,
|
||||
shutdown_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
service_slot: ServiceSlot,
|
||||
}
|
||||
|
||||
impl KnowledgeMcpServer {
|
||||
/// Bind a fresh `127.0.0.1:0` listener, mint a random bearer token, and
|
||||
/// start serving `POST /tool`. The service must be wired separately via
|
||||
/// [`set_service`](Self::set_service) before the first tool call arrives.
|
||||
pub async fn start() -> Result<Self, String> {
|
||||
let auth_token = generate_id();
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.map_err(|e| format!("Failed to bind knowledge MCP HTTP listener: {e}"))?;
|
||||
let http_addr = listener
|
||||
.local_addr()
|
||||
.map_err(|e| format!("Failed to read knowledge MCP local addr: {e}"))?;
|
||||
|
||||
let service_slot: ServiceSlot = Arc::new(RwLock::new(Weak::new()));
|
||||
|
||||
let state = KbMcpState {
|
||||
auth_token: auth_token.clone(),
|
||||
service: service_slot.clone(),
|
||||
};
|
||||
|
||||
let app = axum::Router::new()
|
||||
.route("/tool", axum::routing::post(handle_tool_request))
|
||||
.with_state(state);
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
if let Err(e) = axum::serve(listener, app).await {
|
||||
warn!(error = %e, "Knowledge MCP axum server exited with error");
|
||||
}
|
||||
});
|
||||
|
||||
debug!(http_port = http_addr.port(), "Knowledge MCP Server started (axum)");
|
||||
|
||||
Ok(Self {
|
||||
http_addr,
|
||||
auth_token,
|
||||
shutdown_handle: Some(handle),
|
||||
service_slot,
|
||||
})
|
||||
}
|
||||
|
||||
/// Wire the singleton `KnowledgeService` after it is constructed. Must be
|
||||
/// called once before the first tool request arrives. Takes the `Arc` and
|
||||
/// downgrades internally so callers never construct the `Weak` themselves.
|
||||
pub async fn set_service(&self, svc: &Arc<KnowledgeService>) {
|
||||
// Async setter: the slot is a `tokio::sync::RwLock` (read with
|
||||
// `.read().await` in the handler), so we acquire it with `.write().await`.
|
||||
// `blocking_write` would PANIC here — `set_service` is called from the
|
||||
// async service bootstrap (`AppServices::from_config`), and blocking a
|
||||
// tokio runtime thread is forbidden. Runs once at wiring time, before any
|
||||
// request can contend the slot.
|
||||
*self.service_slot.write().await = Arc::downgrade(svc);
|
||||
}
|
||||
|
||||
pub fn http_port(&self) -> u16 {
|
||||
self.http_addr.port()
|
||||
}
|
||||
|
||||
pub fn auth_token(&self) -> &str {
|
||||
&self.auth_token
|
||||
}
|
||||
|
||||
pub fn stop(&mut self) {
|
||||
if let Some(handle) = self.shutdown_handle.take() {
|
||||
handle.abort();
|
||||
debug!(http_port = self.http_addr.port(), "Knowledge MCP Server stop requested");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for KnowledgeMcpServer {
|
||||
fn drop(&mut self) {
|
||||
self.stop();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Axum handler
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn handle_tool_request(
|
||||
State(state): State<KbMcpState>,
|
||||
headers: axum::http::HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> impl IntoResponse {
|
||||
let provided_token = headers
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.unwrap_or("");
|
||||
|
||||
if provided_token != state.auth_token {
|
||||
warn!("Knowledge MCP: unauthorized request");
|
||||
return (StatusCode::UNAUTHORIZED, Json(json!({"error": "unauthorized"}))).into_response();
|
||||
}
|
||||
|
||||
let tool = body.get("tool").and_then(Value::as_str).unwrap_or("");
|
||||
|
||||
let Some(service) = state.service.read().await.upgrade() else {
|
||||
warn!("Knowledge MCP: service not available");
|
||||
return finish(json!({"error": "knowledge service unavailable"}));
|
||||
};
|
||||
|
||||
// Back-compat: an old bridge may still bake explicit kb_ids. Otherwise scope
|
||||
// is resolved server-side from cwd. Security invariant (all tools): the model
|
||||
// supplies only query/handle/content; scope + write policy are decided
|
||||
// server-side and cannot be widened by the model.
|
||||
let explicit_kb_ids: Vec<String> = body
|
||||
.get("kb_ids")
|
||||
.and_then(|v| serde_json::from_value::<Vec<String>>(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let cwd = body.get("cwd").and_then(Value::as_str).unwrap_or("").to_string();
|
||||
let args = body.get("args").cloned().unwrap_or(Value::Null);
|
||||
|
||||
match tool {
|
||||
"knowledge_search" => {
|
||||
let query = args.get("query").and_then(|q| q.as_str()).unwrap_or("").trim().to_string();
|
||||
let limit = args
|
||||
.get("limit")
|
||||
.and_then(|n| n.as_u64())
|
||||
.map(|n| n as usize)
|
||||
.unwrap_or(8)
|
||||
.clamp(1, 20);
|
||||
let kb_ids = if !explicit_kb_ids.is_empty() {
|
||||
explicit_kb_ids
|
||||
} else {
|
||||
service.resolve_kb_ids_for_cwd(&cwd).await
|
||||
};
|
||||
info!(tool, kb_ids = kb_ids.len(), cwd = %cwd, "Knowledge MCP: dispatching tool");
|
||||
finish(dispatch_search(&service, &kb_ids, &query, limit).await)
|
||||
}
|
||||
"knowledge_read" => {
|
||||
let handle = args.get("handle").and_then(Value::as_str).unwrap_or("").trim().to_string();
|
||||
let kb_ids = if !explicit_kb_ids.is_empty() {
|
||||
explicit_kb_ids
|
||||
} else {
|
||||
service.resolve_kb_ids_for_cwd(&cwd).await
|
||||
};
|
||||
info!(tool, kb_ids = kb_ids.len(), cwd = %cwd, "Knowledge MCP: dispatching tool");
|
||||
finish(dispatch_read(&service, &kb_ids, &handle).await)
|
||||
}
|
||||
"knowledge_write" => {
|
||||
let (bound_kb_ids, binding, wp_key) = service.resolve_write_context_for_cwd(&cwd).await;
|
||||
// Staged inbox scope: prefer an explicit conversation id (per-session
|
||||
// inbox, matching the nomi engine) when the bridge forwards one;
|
||||
// otherwise fall back to the workpath key (per-workspace inbox).
|
||||
let scope = body
|
||||
.get("conversation_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_owned)
|
||||
.unwrap_or(wp_key);
|
||||
info!(tool, kb_ids = bound_kb_ids.len(), cwd = %cwd, "Knowledge MCP: dispatching tool");
|
||||
finish(dispatch_write(&service, &bound_kb_ids, &binding, &scope, &args).await)
|
||||
}
|
||||
_ => {
|
||||
warn!(tool, "Knowledge MCP: unknown tool");
|
||||
finish(json!({"error": format!("unknown tool: {tool}")}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap a JSON body as a response and ask the client to close the connection
|
||||
/// (the stdio bridge runs with `pool_max_idle_per_host(0)` and does not reuse).
|
||||
fn finish(body: Value) -> axum::response::Response {
|
||||
let mut resp = Json(body).into_response();
|
||||
resp.headers_mut()
|
||||
.insert(header::CONNECTION, HeaderValue::from_static("close"));
|
||||
resp
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool implementation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Testable dispatch core: run `search_bases` and render the result envelope.
|
||||
/// Returns `{"result": …}` on success / `{"error": …}` on failure, matching the
|
||||
/// requirement server's envelope.
|
||||
pub(crate) async fn dispatch_search(
|
||||
service: &KnowledgeService,
|
||||
kb_ids: &[String],
|
||||
query: &str,
|
||||
limit: usize,
|
||||
) -> serde_json::Value {
|
||||
match service.search_bases(kb_ids, query, limit).await {
|
||||
Ok(hits) => serde_json::json!({ "result": render_hits(query, &hits) }),
|
||||
Err(e) => serde_json::json!({ "error": e.to_string() }),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a full document by opaque `handle`, scoped to `kb_ids`. A handle whose
|
||||
/// kb_id is outside the resolved scope is rejected — the model cannot widen it.
|
||||
pub(crate) async fn dispatch_read(service: &KnowledgeService, kb_ids: &[String], handle: &str) -> Value {
|
||||
let Some((kb_id, rel_path)) = decode_doc_handle(handle) else {
|
||||
return json!({ "error": format!("invalid document handle: {handle}") });
|
||||
};
|
||||
if !kb_ids.iter().any(|b| b == &kb_id) {
|
||||
return json!({ "error": "handle points to a base not in scope" });
|
||||
}
|
||||
match service.read_file(&kb_id, &rel_path).await {
|
||||
Ok(content) => json!({ "result": content.content }),
|
||||
Err(e) => json!({ "error": e.to_string() }),
|
||||
}
|
||||
}
|
||||
|
||||
/// Write a document through the canonical `write_document` path. The surface is
|
||||
/// always `TerminalAcp` (this server serves ACP/terminal CLIs); the placement
|
||||
/// policy is resolved server-side from the caller's workpath binding — the model
|
||||
/// supplies only `handle | base+rel_path` + `content`, never the policy/scope.
|
||||
pub(crate) async fn dispatch_write(
|
||||
service: &KnowledgeService,
|
||||
bound_kb_ids: &[String],
|
||||
binding: &KnowledgeBinding,
|
||||
scope: &str,
|
||||
args: &Value,
|
||||
) -> Value {
|
||||
let Some(content) = args.get("content").and_then(Value::as_str) else {
|
||||
return json!({ "error": "missing required field: content" });
|
||||
};
|
||||
if content.trim().is_empty() {
|
||||
return json!({ "error": "content is empty" });
|
||||
}
|
||||
let spec = if let Some(handle) = args.get("handle").and_then(Value::as_str).map(str::trim).filter(|s| !s.is_empty()) {
|
||||
WriteTargetSpec::Handle(handle.to_owned())
|
||||
} else {
|
||||
let Some(rel_path) = args.get("rel_path").and_then(Value::as_str).map(str::trim).filter(|s| !s.is_empty()) else {
|
||||
return json!({ "error": "pass either `handle` (to update) or `rel_path` (to create a new document)" });
|
||||
};
|
||||
let kb_id = match resolve_base_id(service, bound_kb_ids, args.get("base").and_then(Value::as_str)).await {
|
||||
Ok(id) => id,
|
||||
Err(e) => return json!({ "error": e }),
|
||||
};
|
||||
WriteTargetSpec::Path { kb_id, rel_path: rel_path.to_owned() }
|
||||
};
|
||||
let policy = resolve_write_policy(WriteSurface::TerminalAcp, binding, scope);
|
||||
let req = WriteRequest { spec, content: content.to_owned(), policy, bound_kb_ids: bound_kb_ids.to_vec() };
|
||||
match service.write_document(req).await {
|
||||
Ok(out) => json!({ "result": {
|
||||
"kb_id": out.kb_id,
|
||||
"rel_path": out.final_rel_path,
|
||||
"staged": out.staged,
|
||||
"updated": matches!(out.op, WriteOp::Update),
|
||||
}}),
|
||||
Err(e) => json!({ "error": e.to_string() }),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a model-supplied base NAME to a bound kb_id (create path). When
|
||||
/// `requested` is omitted and exactly one base is in scope, that base is used.
|
||||
async fn resolve_base_id(service: &KnowledgeService, bound_kb_ids: &[String], requested: Option<&str>) -> Result<String, String> {
|
||||
let bases: Vec<(String, String)> = service
|
||||
.list_bases()
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.filter(|b| bound_kb_ids.contains(&b.id))
|
||||
.map(|b| (b.id, b.name))
|
||||
.collect();
|
||||
if bases.is_empty() {
|
||||
return Err("no knowledge bases are in scope to write to".to_owned());
|
||||
}
|
||||
match requested.map(str::trim).filter(|s| !s.is_empty()) {
|
||||
Some(name) => bases
|
||||
.iter()
|
||||
.find(|(_, n)| n.trim().eq_ignore_ascii_case(name))
|
||||
.map(|(id, _)| id.clone())
|
||||
.ok_or_else(|| {
|
||||
let names = bases.iter().map(|(_, n)| n.as_str()).collect::<Vec<_>>().join(", ");
|
||||
format!("unknown base \"{name}\"; in scope: {names}")
|
||||
}),
|
||||
None => {
|
||||
if bases.len() == 1 {
|
||||
Ok(bases[0].0.clone())
|
||||
} else {
|
||||
let names = bases.iter().map(|(_, n)| n.as_str()).collect::<Vec<_>>().join(", ");
|
||||
Err(format!("multiple bases in scope ({names}); specify `base`"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Render hits into the agent-facing plain-text block the tool returns.
|
||||
fn render_hits(query: &str, hits: &[KnowledgeSearchHit]) -> String {
|
||||
if hits.is_empty() {
|
||||
return format!("No matches for \"{query}\" in the mounted knowledge bases. Try different terms.");
|
||||
}
|
||||
let mut out = format!("{} result(s) for \"{}\":\n", hits.len(), query);
|
||||
for (i, h) in hits.iter().enumerate() {
|
||||
out.push_str(&format!(
|
||||
"{}. [{}] {} — {}\n {}\n handle: {}\n",
|
||||
i + 1,
|
||||
h.kb_name,
|
||||
h.rel_path,
|
||||
if h.heading.is_empty() { "(no heading)" } else { &h.heading },
|
||||
h.snippet,
|
||||
encode_doc_handle(&h.kb_id, &h.rel_path),
|
||||
));
|
||||
}
|
||||
out.push_str(
|
||||
"\nTo read a full document, call knowledge_read with its `handle`. To update one, call \
|
||||
knowledge_write with that same `handle` (do NOT rebuild the path).",
|
||||
);
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::events::KnowledgeEventEmitter;
|
||||
|
||||
#[derive(Default)]
|
||||
struct NoopBroadcaster;
|
||||
impl nomifun_realtime::EventBroadcaster for NoopBroadcaster {
|
||||
fn broadcast(&self, _event: nomifun_api_types::WebSocketMessage<serde_json::Value>) {}
|
||||
}
|
||||
|
||||
fn hit(kb_name: &str, rel_path: &str, heading: &str, snippet: &str) -> KnowledgeSearchHit {
|
||||
KnowledgeSearchHit {
|
||||
kb_id: "kb_1".into(),
|
||||
kb_name: kb_name.into(),
|
||||
rel_path: rel_path.into(),
|
||||
heading: heading.into(),
|
||||
snippet: snippet.into(),
|
||||
score: 1,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_hits_empty_reports_no_matches() {
|
||||
let out = render_hits("回滚", &[]);
|
||||
assert!(out.contains("No matches"), "got: {out}");
|
||||
assert!(out.contains("回滚"), "echoes the query: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_hits_non_empty_lists_path_heading_and_handle() {
|
||||
let hits = vec![hit("运维手册", "rollback.md", "回滚流程", "回滚分三步")];
|
||||
let out = render_hits("回滚", &hits);
|
||||
assert!(out.contains("rollback.md"), "path: {out}");
|
||||
assert!(out.contains("回滚流程"), "heading: {out}");
|
||||
assert!(out.contains("运维手册"), "kb name: {out}");
|
||||
assert!(out.contains("handle: kdoc_"), "handle: {out}");
|
||||
assert!(out.contains("knowledge_read") || out.contains("knowledge_write"), "tool hint: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_hits_blank_heading_falls_back() {
|
||||
let hits = vec![hit("库", "a.md", "", "some snippet")];
|
||||
let out = render_hits("topic", &hits);
|
||||
assert!(out.contains("(no heading)"), "got: {out}");
|
||||
}
|
||||
|
||||
/// Build a real `KnowledgeService` over an in-memory DB + temp data dir
|
||||
/// (recipe from nomifun-ai-agent's `knowledge_search_e2e`). Returns the
|
||||
/// service and the `TempDir` (keep it alive for the test's duration).
|
||||
async fn build_service() -> (Arc<KnowledgeService>, tempfile::TempDir) {
|
||||
let db = nomifun_db::init_database_memory().await.expect("in-memory db");
|
||||
let repo = Arc::new(nomifun_db::SqliteKnowledgeRepository::new(db.pool().clone()));
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let emitter = KnowledgeEventEmitter::new(Arc::new(NoopBroadcaster));
|
||||
let svc = Arc::new(KnowledgeService::new(repo, tmp.path(), emitter));
|
||||
(svc, tmp)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_search_finds_doc_and_wraps_result() {
|
||||
let (svc, _tmp) = build_service().await;
|
||||
let info = svc.create_base("运维手册", "", None, None).await.unwrap();
|
||||
let root = svc.data_dir().join("knowledge").join(&info.id);
|
||||
// The self-ignore the mount writes — must NOT blind the search.
|
||||
std::fs::write(root.join(".gitignore"), "*\n").unwrap();
|
||||
std::fs::write(root.join("rollback.md"), "# 回滚流程\n回滚分三步\n").unwrap();
|
||||
|
||||
let out = dispatch_search(&svc, &[info.id], "回滚", 8).await;
|
||||
let result = out
|
||||
.get("result")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_else(|| panic!("expected a result envelope, got {out}"));
|
||||
assert!(result.contains("rollback.md"), "must surface the doc:\n{result}");
|
||||
assert!(result.contains("回滚流程"), "must include heading:\n{result}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_search_no_match_reports_cleanly() {
|
||||
let (svc, _tmp) = build_service().await;
|
||||
let info = svc.create_base("库", "", None, None).await.unwrap();
|
||||
let root = svc.data_dir().join("knowledge").join(&info.id);
|
||||
std::fs::write(root.join("a.md"), "# A\nunrelated content\n").unwrap();
|
||||
|
||||
let out = dispatch_search(&svc, &[info.id], "完全不存在的主题词", 8).await;
|
||||
let result = out.get("result").and_then(Value::as_str).unwrap_or_else(|| panic!("got {out}"));
|
||||
assert!(result.contains("No matches"), "got: {result}");
|
||||
}
|
||||
|
||||
// ── cwd-based scope resolution (Task 5) ─────────────────────────────
|
||||
|
||||
/// Helper: start a `KnowledgeMcpServer`, wire a service, and return
|
||||
/// (server, service, port, token) for HTTP-level tests.
|
||||
async fn start_wired_server() -> (KnowledgeMcpServer, Arc<KnowledgeService>, u16, String, tempfile::TempDir) {
|
||||
let (svc, tmp) = build_service().await;
|
||||
let server = KnowledgeMcpServer::start().await.expect("bind");
|
||||
server.set_service(&svc).await;
|
||||
let port = server.http_port();
|
||||
let token = server.auth_token().to_owned();
|
||||
(server, svc, port, token, tmp)
|
||||
}
|
||||
|
||||
/// POST /tool with a JSON body, return the response JSON.
|
||||
async fn post_tool(port: u16, token: &str, body: Value) -> Value {
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client
|
||||
.post(format!("http://127.0.0.1:{port}/tool"))
|
||||
.header("Authorization", format!("Bearer {token}"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request");
|
||||
resp.json::<Value>().await.expect("json")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_request_with_cwd_resolves_scope_via_service() {
|
||||
let (_server, svc, port, token, _tmp) = start_wired_server().await;
|
||||
|
||||
// Create a base and bind it to a workpath.
|
||||
let info = svc.create_base("项目库", "", None, None).await.unwrap();
|
||||
let root = svc.data_dir().join("knowledge").join(&info.id);
|
||||
std::fs::write(root.join("api.md"), "# API\n接口文档内容\n").unwrap();
|
||||
|
||||
let ws = "/Users/test/myproject";
|
||||
let key = crate::workpath::workpath_key(ws);
|
||||
svc.set_binding(
|
||||
crate::workpath::WORKPATH_BINDING_KIND,
|
||||
&key,
|
||||
crate::service::KnowledgeBinding {
|
||||
enabled: true,
|
||||
kb_ids: vec![info.id.clone()],
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Request with cwd (no kb_ids) → uses cwd-resolved scope.
|
||||
let resp = post_tool(port, &token, json!({
|
||||
"tool": "knowledge_search",
|
||||
"cwd": ws,
|
||||
"args": { "query": "接口" }
|
||||
}))
|
||||
.await;
|
||||
let result = resp.get("result").and_then(Value::as_str)
|
||||
.unwrap_or_else(|| panic!("expected result, got {resp}"));
|
||||
assert!(result.contains("api.md"), "cwd scope should find the doc: {result}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_request_with_explicit_kb_ids_uses_them_backcompat() {
|
||||
let (_server, svc, port, token, _tmp) = start_wired_server().await;
|
||||
|
||||
let info = svc.create_base("手册", "", None, None).await.unwrap();
|
||||
let root = svc.data_dir().join("knowledge").join(&info.id);
|
||||
std::fs::write(root.join("ops.md"), "# Ops\n运维流程\n").unwrap();
|
||||
|
||||
// Create another base (not bound to anything).
|
||||
let info2 = svc.create_base("无关库", "", None, None).await.unwrap();
|
||||
let root2 = svc.data_dir().join("knowledge").join(&info2.id);
|
||||
std::fs::write(root2.join("other.md"), "# Other\n别的东西\n").unwrap();
|
||||
|
||||
// Request with explicit kb_ids (old bridge style) → uses those, ignores cwd.
|
||||
let resp = post_tool(port, &token, json!({
|
||||
"tool": "knowledge_search",
|
||||
"kb_ids": [info.id],
|
||||
"cwd": "/some/unbound/path",
|
||||
"args": { "query": "运维" }
|
||||
}))
|
||||
.await;
|
||||
let result = resp.get("result").and_then(Value::as_str)
|
||||
.unwrap_or_else(|| panic!("expected result, got {resp}"));
|
||||
assert!(result.contains("ops.md"), "explicit kb_ids should be used: {result}");
|
||||
// The other base should NOT be searched (explicit kb_ids narrows scope).
|
||||
assert!(!result.contains("other.md"), "should not search unspecified bases: {result}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_request_with_empty_cwd_searches_all_bases() {
|
||||
let (_server, svc, port, token, _tmp) = start_wired_server().await;
|
||||
|
||||
let info = svc.create_base("全局库", "", None, None).await.unwrap();
|
||||
let root = svc.data_dir().join("knowledge").join(&info.id);
|
||||
std::fs::write(root.join("global.md"), "# Global\n全局知识\n").unwrap();
|
||||
|
||||
// No kb_ids, empty cwd → fallback to all bases.
|
||||
let resp = post_tool(port, &token, json!({
|
||||
"tool": "knowledge_search",
|
||||
"cwd": "",
|
||||
"args": { "query": "全局" }
|
||||
}))
|
||||
.await;
|
||||
let result = resp.get("result").and_then(Value::as_str)
|
||||
.unwrap_or_else(|| panic!("expected result, got {resp}"));
|
||||
assert!(result.contains("global.md"), "empty cwd should search all: {result}");
|
||||
}
|
||||
|
||||
// ── knowledge_read / knowledge_write (P2) ───────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_read_returns_content_within_scope_and_denies_outside() {
|
||||
let (svc, _tmp) = build_service().await;
|
||||
let info = svc.create_base("库", "", None, None).await.unwrap();
|
||||
svc.write_file(&info.id, "terms.md", "# T\nBODY-市盈率").await.unwrap();
|
||||
let h = encode_doc_handle(&info.id, "terms.md");
|
||||
|
||||
let ok = dispatch_read(&svc, std::slice::from_ref(&info.id), &h).await;
|
||||
assert!(ok.get("result").and_then(Value::as_str).unwrap_or("").contains("BODY-市盈率"), "{ok}");
|
||||
// Out of scope (empty kb_ids) → denied.
|
||||
let denied = dispatch_read(&svc, &[], &h).await;
|
||||
assert!(denied.get("error").is_some(), "out-of-scope handle must be denied: {denied}");
|
||||
// Malformed handle → error.
|
||||
let bad = dispatch_read(&svc, std::slice::from_ref(&info.id), "not-a-handle").await;
|
||||
assert!(bad.get("error").is_some(), "{bad}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_write_staged_lands_in_inbox_and_preserves_original() {
|
||||
let (svc, _tmp) = build_service().await;
|
||||
let info = svc.create_base("库", "", None, None).await.unwrap();
|
||||
svc.write_file(&info.id, "terms.md", "ORIGINAL").await.unwrap();
|
||||
let binding = KnowledgeBinding {
|
||||
enabled: true,
|
||||
writeback: true,
|
||||
writeback_mode: "staged".into(),
|
||||
kb_ids: vec![info.id.clone()],
|
||||
..Default::default()
|
||||
};
|
||||
let out = dispatch_write(
|
||||
&svc,
|
||||
std::slice::from_ref(&info.id),
|
||||
&binding,
|
||||
"conv-x",
|
||||
&json!({ "handle": encode_doc_handle(&info.id, "terms.md"), "content": "PROPOSED" }),
|
||||
)
|
||||
.await;
|
||||
let r = out.get("result").unwrap_or_else(|| panic!("{out}"));
|
||||
assert_eq!(r.get("rel_path").and_then(Value::as_str), Some("_inbox/conv-x/terms.md"));
|
||||
assert_eq!(r.get("staged").and_then(Value::as_bool), Some(true));
|
||||
// Original untouched; proposal staged.
|
||||
assert_eq!(svc.read_file(&info.id, "terms.md").await.unwrap().content, "ORIGINAL");
|
||||
assert_eq!(svc.read_file(&info.id, "_inbox/conv-x/terms.md").await.unwrap().content, "PROPOSED");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_write_refused_when_writeback_disabled() {
|
||||
let (svc, _tmp) = build_service().await;
|
||||
let info = svc.create_base("库", "", None, None).await.unwrap();
|
||||
svc.write_file(&info.id, "terms.md", "x").await.unwrap();
|
||||
// Binding present but writeback off → policy Disabled.
|
||||
let binding = KnowledgeBinding { enabled: true, writeback: false, kb_ids: vec![info.id.clone()], ..Default::default() };
|
||||
let out = dispatch_write(
|
||||
&svc,
|
||||
std::slice::from_ref(&info.id),
|
||||
&binding,
|
||||
"wp",
|
||||
&json!({ "handle": encode_doc_handle(&info.id, "terms.md"), "content": "y" }),
|
||||
)
|
||||
.await;
|
||||
assert!(out.get("error").is_some(), "writeback off must refuse: {out}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_knowledge_write_routes_through_policy_direct() {
|
||||
let (_server, svc, port, token, _tmp) = start_wired_server().await;
|
||||
let info = svc.create_base("项目库", "", None, None).await.unwrap();
|
||||
svc.write_file(&info.id, "notes.md", "OLD").await.unwrap();
|
||||
let ws = "/Users/test/wp-write";
|
||||
let key = crate::workpath::workpath_key(ws);
|
||||
svc.set_binding(
|
||||
crate::workpath::WORKPATH_BINDING_KIND,
|
||||
&key,
|
||||
KnowledgeBinding {
|
||||
enabled: true,
|
||||
writeback: true,
|
||||
writeback_mode: "direct".into(),
|
||||
kb_ids: vec![info.id.clone()],
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = post_tool(port, &token, json!({
|
||||
"tool": "knowledge_write",
|
||||
"cwd": ws,
|
||||
"args": { "handle": encode_doc_handle(&info.id, "notes.md"), "content": "NEW" }
|
||||
}))
|
||||
.await;
|
||||
assert!(resp.get("result").is_some(), "expected result, got {resp}");
|
||||
assert_eq!(svc.read_file(&info.id, "notes.md").await.unwrap().content, "NEW");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
//! Platform-aware mount engine: materializes knowledge bases inside a
|
||||
//! workspace at `.nomi/knowledge/{link_name}` using NTFS junctions on
|
||||
//! Windows (no privilege required), symlinks on Unix, and a recursive copy
|
||||
//! as last-resort fallback (same degradation strategy as the skill linker in
|
||||
//! `nomifun-extension`).
|
||||
//!
|
||||
//! The mount directory is wholly owned by this module: anything inside it
|
||||
//! that is not in the desired set (or in [`MANAGED_KEEP`]) gets removed on
|
||||
//! the next sync. Targets are never touched — removal only deletes the link
|
||||
//! (or the fallback copy). Sibling `.nomi/` trees (`.nomi/skills`, …) are
|
||||
//! never touched either.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::{KB_LEGACY_MOUNT_REL_DIR, KB_MOUNT_REL_DIR};
|
||||
|
||||
/// One desired mount: `{workspace}/.nomi/knowledge/{link_name}` → `target`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MountSpec {
|
||||
pub link_name: String,
|
||||
pub target: PathBuf,
|
||||
}
|
||||
|
||||
/// Platform-managed companion files inside the mount root (the self-ignore,
|
||||
/// the terminal-facing README) — exempt from the stale-entry sweep, and
|
||||
/// reserved against base link names (see `service::unique_link_name`).
|
||||
pub(crate) const MANAGED_KEEP: &[&str] = &[".gitignore", "README.md"];
|
||||
|
||||
/// Synchronize the workspace mount directory to exactly `specs`.
|
||||
///
|
||||
/// Returns the link names that are present (linked or copied) after the
|
||||
/// sync. Individual failures are logged and skipped — mounting must never
|
||||
/// brick a session start.
|
||||
pub async fn sync_mounts(workspace: &Path, specs: Vec<MountSpec>) -> Vec<String> {
|
||||
let workspace = workspace.to_path_buf();
|
||||
tokio::task::spawn_blocking(move || sync_mounts_blocking(&workspace, &specs))
|
||||
.await
|
||||
.unwrap_or_else(|e| {
|
||||
tracing::warn!(error = %e, "knowledge mount task join error");
|
||||
Vec::new()
|
||||
})
|
||||
}
|
||||
|
||||
fn sync_mounts_blocking(workspace: &Path, specs: &[MountSpec]) -> Vec<String> {
|
||||
let present = sync_mounts_inner(workspace, specs);
|
||||
cleanup_legacy_mount_root(workspace);
|
||||
present
|
||||
}
|
||||
|
||||
fn sync_mounts_inner(workspace: &Path, specs: &[MountSpec]) -> Vec<String> {
|
||||
let mount_root = workspace.join(KB_MOUNT_REL_DIR);
|
||||
|
||||
if specs.is_empty() {
|
||||
// Nothing should be mounted: clear our directory if it exists, then
|
||||
// try to remove the (now empty) scaffolding. The parent `.nomi/` is
|
||||
// only removed when empty — sibling trees keep it alive. Errors are
|
||||
// non-fatal.
|
||||
if mount_root.exists() {
|
||||
if let Ok(entries) = std::fs::read_dir(&mount_root) {
|
||||
for entry in entries.flatten() {
|
||||
remove_mount_entry(&entry.path());
|
||||
}
|
||||
}
|
||||
let _ = std::fs::remove_dir(&mount_root);
|
||||
if let Some(parent) = mount_root.parent() {
|
||||
let _ = std::fs::remove_dir(parent);
|
||||
}
|
||||
}
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
if let Err(e) = std::fs::create_dir_all(&mount_root) {
|
||||
tracing::warn!(path = %mount_root.display(), error = %e, "failed to create knowledge mount dir");
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Self-ignore the mount directory: when the workspace is a user git
|
||||
// repo, junctions would otherwise expose the knowledge base content as
|
||||
// committable project files. The ignore file lives INSIDE
|
||||
// `.nomi/knowledge/` — never at the `.nomi/` root — so committable
|
||||
// siblings like `.nomi/skills` stay visible to git.
|
||||
let gitignore = mount_root.join(".gitignore");
|
||||
if !gitignore.exists() {
|
||||
if let Err(e) = std::fs::write(&gitignore, "*\n") {
|
||||
tracing::warn!(path = %gitignore.display(), error = %e, "failed to write knowledge mount .gitignore");
|
||||
}
|
||||
}
|
||||
|
||||
let desired: HashMap<&str, &MountSpec> = specs.iter().map(|s| (s.link_name.as_str(), s)).collect();
|
||||
|
||||
// Pass 1: drop stale entries and stale links whose target changed.
|
||||
if let Ok(entries) = std::fs::read_dir(&mount_root) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if MANAGED_KEEP.contains(&name.as_str()) {
|
||||
continue;
|
||||
}
|
||||
match desired.get(name.as_str()) {
|
||||
None => remove_mount_entry(&path),
|
||||
Some(spec) => {
|
||||
if let Some(current) = read_link_target(&path)
|
||||
&& current != spec.target
|
||||
{
|
||||
remove_mount_entry(&path);
|
||||
}
|
||||
// A non-link entry (copy fallback) is left in place: we
|
||||
// cannot cheaply verify it, and re-copying every session
|
||||
// start would be wasteful. It gets refreshed whenever the
|
||||
// base set changes its name (different link_name).
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pass 2: create whatever is missing.
|
||||
let mut present = Vec::new();
|
||||
for spec in specs {
|
||||
let link = mount_root.join(&spec.link_name);
|
||||
if link.exists() || read_link_target(&link).is_some() {
|
||||
present.push(spec.link_name.clone());
|
||||
continue;
|
||||
}
|
||||
if !spec.target.is_dir() {
|
||||
tracing::warn!(
|
||||
target = %spec.target.display(),
|
||||
name = %spec.link_name,
|
||||
"knowledge base root missing; skipping mount"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
match create_link(&spec.target, &link) {
|
||||
Ok(()) => present.push(spec.link_name.clone()),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
target = %spec.target.display(),
|
||||
link = %link.display(),
|
||||
error = %e,
|
||||
raw_os_error = ?e.raw_os_error(),
|
||||
"knowledge link failed; falling back to copy"
|
||||
);
|
||||
match copy_dir_recursive(&spec.target, &link) {
|
||||
Ok(()) => present.push(spec.link_name.clone()),
|
||||
Err(e) => {
|
||||
tracing::warn!(link = %link.display(), error = %e, "knowledge copy fallback failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
present
|
||||
}
|
||||
|
||||
/// Best-effort sweep of the pre-`.nomi` mount scaffolding left in workspaces
|
||||
/// created before the rename: `{ws}/.nomifun/knowledge/*` links (deleted as
|
||||
/// links — never followed into the knowledge bases) plus the self-ignore we
|
||||
/// used to write at `{ws}/.nomifun/.gitignore`. The `.nomifun/` directory
|
||||
/// itself is only removed when empty, so unrelated user files keep both the
|
||||
/// file and the directory alive. Idempotent; failures only warn.
|
||||
fn cleanup_legacy_mount_root(workspace: &Path) {
|
||||
let legacy_mount = workspace.join(KB_LEGACY_MOUNT_REL_DIR);
|
||||
if legacy_mount.exists() {
|
||||
if let Ok(entries) = std::fs::read_dir(&legacy_mount) {
|
||||
for entry in entries.flatten() {
|
||||
remove_mount_entry(&entry.path());
|
||||
}
|
||||
}
|
||||
if let Err(e) = std::fs::remove_dir(&legacy_mount) {
|
||||
tracing::warn!(path = %legacy_mount.display(), error = %e, "failed to remove legacy knowledge mount dir");
|
||||
}
|
||||
}
|
||||
let Some(legacy_root) = legacy_mount.parent() else { return };
|
||||
let gitignore = legacy_root.join(".gitignore");
|
||||
// Only delete the ignore file we wrote (content `*`) — anything else is
|
||||
// the user's and stays.
|
||||
if matches!(std::fs::read_to_string(&gitignore), Ok(content) if content.trim() == "*") {
|
||||
if let Err(e) = std::fs::remove_file(&gitignore) {
|
||||
tracing::warn!(path = %gitignore.display(), error = %e, "failed to remove legacy mount .gitignore");
|
||||
}
|
||||
}
|
||||
let _ = std::fs::remove_dir(legacy_root);
|
||||
}
|
||||
|
||||
/// Remove one entry inside the mount dir without ever touching the link
|
||||
/// target's contents: junctions/symlinks are removed as links; plain
|
||||
/// directories (copy fallback leftovers) are removed recursively — they are
|
||||
/// copies we created, never user originals.
|
||||
fn remove_mount_entry(path: &Path) {
|
||||
let result = if read_link_target(path).is_some() {
|
||||
if path.is_dir() {
|
||||
std::fs::remove_dir(path)
|
||||
} else {
|
||||
std::fs::remove_file(path)
|
||||
}
|
||||
} else if path.is_dir() {
|
||||
std::fs::remove_dir_all(path)
|
||||
} else {
|
||||
std::fs::remove_file(path)
|
||||
};
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(path = %path.display(), error = %e, "failed to remove stale knowledge mount entry");
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the target of a symlink or (on Windows) NTFS junction; `None` for
|
||||
/// regular files/dirs or when the entry does not exist.
|
||||
fn read_link_target(path: &Path) -> Option<PathBuf> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
if junction::exists(path).unwrap_or(false) {
|
||||
return junction::get_target(path).ok();
|
||||
}
|
||||
}
|
||||
let meta = std::fs::symlink_metadata(path).ok()?;
|
||||
if meta.file_type().is_symlink() {
|
||||
std::fs::read_link(path).ok()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn create_link(src: &Path, dst: &Path) -> io::Result<()> {
|
||||
std::os::unix::fs::symlink(src, dst)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn create_link(src: &Path, dst: &Path) -> io::Result<()> {
|
||||
// Junctions work without SeCreateSymbolicLink (Developer Mode/Admin),
|
||||
// which most users don't have — mirrors the skill linker's rationale.
|
||||
junction::create(src, dst)
|
||||
}
|
||||
|
||||
fn copy_dir_recursive(src: &Path, dst: &Path) -> io::Result<()> {
|
||||
std::fs::create_dir_all(dst)?;
|
||||
for entry in walkdir::WalkDir::new(src).min_depth(1) {
|
||||
let entry = entry.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
|
||||
let rel = entry
|
||||
.path()
|
||||
.strip_prefix(src)
|
||||
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
|
||||
let to = dst.join(rel);
|
||||
if entry.file_type().is_dir() {
|
||||
std::fs::create_dir_all(&to)?;
|
||||
} else if entry.file_type().is_file() {
|
||||
if let Some(parent) = to.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::copy(entry.path(), &to)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn make_base(dir: &TempDir, name: &str) -> PathBuf {
|
||||
let root = dir.path().join(name);
|
||||
std::fs::create_dir_all(&root).unwrap();
|
||||
std::fs::write(root.join("note.md"), "# hi").unwrap();
|
||||
root
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mounts_link_and_cleanup() {
|
||||
let bases = TempDir::new().unwrap();
|
||||
let ws = TempDir::new().unwrap();
|
||||
let kb_a = make_base(&bases, "kb_a");
|
||||
let kb_b = make_base(&bases, "kb_b");
|
||||
|
||||
// Mount both.
|
||||
let present = sync_mounts(
|
||||
ws.path(),
|
||||
vec![
|
||||
MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb_a.clone(),
|
||||
},
|
||||
MountSpec {
|
||||
link_name: "乙".into(),
|
||||
target: kb_b.clone(),
|
||||
},
|
||||
],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(present.len(), 2);
|
||||
let mount_root = ws.path().join(KB_MOUNT_REL_DIR);
|
||||
assert!(mount_root.join("甲").join("note.md").exists());
|
||||
assert!(mount_root.join("乙").join("note.md").exists());
|
||||
// The mount dir self-ignores so junction content never leaks into
|
||||
// the user's git repository.
|
||||
let gitignore = mount_root.join(".gitignore");
|
||||
assert_eq!(std::fs::read_to_string(&gitignore).unwrap().trim(), "*");
|
||||
|
||||
// Shrink to one — the other must disappear, target stays intact.
|
||||
let present = sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb_a.clone(),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(present, vec!["甲".to_string()]);
|
||||
assert!(!mount_root.join("乙").exists());
|
||||
assert!(kb_b.join("note.md").exists(), "unmount must not delete target content");
|
||||
|
||||
// Retarget the same name — link must follow.
|
||||
let present = sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb_b.clone(),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(present.len(), 1);
|
||||
std::fs::write(kb_b.join("only_b.md"), "b").unwrap();
|
||||
assert!(mount_root.join("甲").join("only_b.md").exists());
|
||||
|
||||
// Empty set clears the scaffolding.
|
||||
let present = sync_mounts(ws.path(), vec![]).await;
|
||||
assert!(present.is_empty());
|
||||
assert!(!mount_root.exists());
|
||||
assert!(kb_a.join("note.md").exists());
|
||||
assert!(kb_b.join("note.md").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gitignore_written_inside_knowledge_dir() {
|
||||
let bases = TempDir::new().unwrap();
|
||||
let ws = TempDir::new().unwrap();
|
||||
let kb = make_base(&bases, "kb_g");
|
||||
|
||||
sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb,
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
|
||||
// The self-ignore lives INSIDE `.nomi/knowledge/` — pinned to the
|
||||
// literal path so a constant regression cannot slip through.
|
||||
let inside = ws.path().join(".nomi").join("knowledge").join(".gitignore");
|
||||
assert_eq!(std::fs::read_to_string(&inside).unwrap().trim(), "*");
|
||||
// Never at the `.nomi/` root: that would shadow committable sibling
|
||||
// trees like `.nomi/skills` out of the user's git repository.
|
||||
assert!(!ws.path().join(".nomi").join(".gitignore").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_nomifun_mounts_cleaned() {
|
||||
let bases = TempDir::new().unwrap();
|
||||
let ws = TempDir::new().unwrap();
|
||||
let kb = make_base(&bases, "kb_legacy");
|
||||
|
||||
// Scenario 1: `.nomifun/` holds only our scaffolding (a mounted link
|
||||
// + the self-ignore) → the whole legacy dir disappears.
|
||||
let legacy_root = ws.path().join(".nomifun");
|
||||
let legacy_knowledge = legacy_root.join("knowledge");
|
||||
std::fs::create_dir_all(&legacy_knowledge).unwrap();
|
||||
std::fs::write(legacy_root.join(".gitignore"), "*\n").unwrap();
|
||||
let legacy_link = legacy_knowledge.join("旧库");
|
||||
if create_link(&kb, &legacy_link).is_err() {
|
||||
// Platform refused the link (CI sandbox): a plain dir still
|
||||
// exercises the cleanup path (copy-fallback leftovers are dirs).
|
||||
std::fs::create_dir_all(&legacy_link).unwrap();
|
||||
}
|
||||
|
||||
sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb.clone(),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(!legacy_root.exists(), "legacy .nomifun scaffolding must be fully removed");
|
||||
assert!(
|
||||
kb.join("note.md").exists(),
|
||||
"legacy cleanup must delete links as links, never follow into the base"
|
||||
);
|
||||
|
||||
// Scenario 2: `.nomifun/` also holds an unrelated user file → only
|
||||
// our pieces go; the directory and the user file survive.
|
||||
std::fs::create_dir_all(&legacy_knowledge).unwrap();
|
||||
std::fs::write(legacy_root.join(".gitignore"), "*\n").unwrap();
|
||||
std::fs::write(legacy_root.join("user-note.txt"), "keep me").unwrap();
|
||||
|
||||
sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb.clone(),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(!legacy_knowledge.exists());
|
||||
assert!(!legacy_root.join(".gitignore").exists());
|
||||
assert_eq!(std::fs::read_to_string(legacy_root.join("user-note.txt")).unwrap(), "keep me");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn managed_files_survive_sync() {
|
||||
let bases = TempDir::new().unwrap();
|
||||
let ws = TempDir::new().unwrap();
|
||||
let kb = make_base(&bases, "kb_m");
|
||||
let spec = || {
|
||||
vec![MountSpec {
|
||||
link_name: "甲".into(),
|
||||
target: kb.clone(),
|
||||
}]
|
||||
};
|
||||
|
||||
sync_mounts(ws.path(), spec()).await;
|
||||
let mount_root = ws.path().join(KB_MOUNT_REL_DIR);
|
||||
// Platform-managed companion file (terminal README, see MANAGED_KEEP)
|
||||
// must not be swept as a stale mount on the next sync.
|
||||
std::fs::write(mount_root.join("README.md"), "# managed").unwrap();
|
||||
|
||||
sync_mounts(ws.path(), spec()).await;
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(mount_root.join("README.md")).unwrap(),
|
||||
"# managed"
|
||||
);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(mount_root.join(".gitignore")).unwrap().trim(),
|
||||
"*"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_target_is_skipped() {
|
||||
let ws = TempDir::new().unwrap();
|
||||
let present = sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "ghost".into(),
|
||||
target: PathBuf::from("Z:/definitely/not/here"),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
assert!(present.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn writes_through_mount_reach_target() {
|
||||
let bases = TempDir::new().unwrap();
|
||||
let ws = TempDir::new().unwrap();
|
||||
let kb = make_base(&bases, "kb_w");
|
||||
|
||||
sync_mounts(
|
||||
ws.path(),
|
||||
vec![MountSpec {
|
||||
link_name: "w".into(),
|
||||
target: kb.clone(),
|
||||
}],
|
||||
)
|
||||
.await;
|
||||
|
||||
let mounted = ws.path().join(KB_MOUNT_REL_DIR).join("w");
|
||||
// Skip the assertion when the platform degraded to a copy (no link
|
||||
// semantics) — detectable because read_link_target returns None.
|
||||
if read_link_target(&mounted).is_some() {
|
||||
std::fs::write(mounted.join("written.md"), "wb").unwrap();
|
||||
assert!(kb.join("written.md").exists(), "write-back must land in the base root");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,900 @@
|
||||
//! `/api/knowledge/*` route handlers.
|
||||
|
||||
use axum::Router;
|
||||
use axum::extract::rejection::JsonRejection;
|
||||
use axum::extract::{Extension, Json, Path, Query, State};
|
||||
use axum::routing::{get, post};
|
||||
|
||||
use nomifun_api_types::{ApiResponse, ConnectorCredentialSummary, CreateKnowledgeTagRequest, KnowledgeSource, KnowledgeTag, UpdateKnowledgeTagRequest};
|
||||
use nomifun_auth::CurrentUser;
|
||||
use nomifun_common::AppError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::connector::ConnectorIdentity;
|
||||
use crate::export::{self, ExportSummary, ImportSummary};
|
||||
use crate::service::{
|
||||
AutogenOutcome, ConsumerInfo, InboxDiff, InboxEntry, InboxMergeResult, KbFileContent, KbFileEntry,
|
||||
KnowledgeBaseInfo, KnowledgeBinding, KnowledgeSearchHit, RefreshSourceSummary,
|
||||
};
|
||||
use crate::state::KnowledgeRouterState;
|
||||
|
||||
pub fn knowledge_routes(state: KnowledgeRouterState) -> Router {
|
||||
Router::new()
|
||||
.route("/api/knowledge/bases", get(list_bases).post(create_base))
|
||||
.route("/api/knowledge/bases/import", post(import_base))
|
||||
.route(
|
||||
"/api/knowledge/bases/{id}",
|
||||
get(get_base).put(update_base).delete(delete_base),
|
||||
)
|
||||
.route("/api/knowledge/bases/{id}/export", post(export_base))
|
||||
.route("/api/knowledge/bases/{id}/autogen", post(autogen_base))
|
||||
.route("/api/knowledge/description/generate", post(generate_description))
|
||||
.route("/api/knowledge/description/polish", post(polish_description))
|
||||
.route("/api/knowledge/bases/{id}/refresh-source", post(refresh_source))
|
||||
.route("/api/knowledge/bases/{id}/source", axum::routing::put(set_source))
|
||||
.route("/api/knowledge/bases/{id}/sync", post(sync_source))
|
||||
.route(
|
||||
"/api/knowledge/connectors/credentials",
|
||||
get(list_credentials).post(create_credential),
|
||||
)
|
||||
.route(
|
||||
"/api/knowledge/connectors/credentials/{id}",
|
||||
axum::routing::delete(delete_credential),
|
||||
)
|
||||
.route(
|
||||
"/api/knowledge/connectors/credentials/{id}/test",
|
||||
post(test_credential),
|
||||
)
|
||||
.route(
|
||||
"/api/knowledge/tags",
|
||||
get(list_tags).post(create_tag),
|
||||
)
|
||||
.route(
|
||||
"/api/knowledge/tags/{key}",
|
||||
axum::routing::put(update_tag).delete(delete_tag),
|
||||
)
|
||||
.route("/api/knowledge/bases/{id}/files", get(list_files))
|
||||
.route("/api/knowledge/bases/{id}/inbox", get(list_inbox))
|
||||
.route("/api/knowledge/inbox/pending-count", get(pending_inbox_count))
|
||||
.route("/api/knowledge/bases/{id}/inbox/diff", get(inbox_diff))
|
||||
.route("/api/knowledge/bases/{id}/inbox/merge", post(merge_inbox))
|
||||
.route("/api/knowledge/bases/{id}/inbox/discard", post(discard_inbox))
|
||||
.route("/api/knowledge/inbox/merge-all", post(merge_all_inbox))
|
||||
.route("/api/knowledge/inbox/discard-all", post(discard_all_inbox))
|
||||
.route("/api/knowledge/bases/{id}/consumers", get(list_consumers))
|
||||
.route(
|
||||
"/api/knowledge/bases/{id}/file",
|
||||
get(read_file).put(write_file).delete(delete_file),
|
||||
)
|
||||
.route(
|
||||
// `target_id` is ONE path segment. Workpath targets (normalized
|
||||
// absolute paths) therefore arrive percent-encoded — the
|
||||
// frontend calls `encodeURIComponent(workpathKey)` so `/`
|
||||
// travels as `%2F`. axum matches routes on the still-encoded
|
||||
// path and the `Path` extractor decodes afterwards, so an
|
||||
// encoded path never splits into extra segments (pinned by
|
||||
// `binding_route_extracts_percent_encoded_workpath` below).
|
||||
"/api/knowledge/binding/{kind}/{target_id}",
|
||||
get(get_binding).post(set_binding),
|
||||
)
|
||||
.route("/api/knowledge/search", post(search_bases))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn list_bases(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
) -> Result<Json<ApiResponse<Vec<KnowledgeBaseInfo>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_bases().await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateBaseRequest {
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
/// Absolute path of an existing external directory; omit to provision a
|
||||
/// managed directory under the backend data dir.
|
||||
root_path: Option<String>,
|
||||
/// Optional URL source, stored in `extra.source`. `mode=live` stores it
|
||||
/// without fetching; `mode=snapshot` fetches every entry into
|
||||
/// `snapshots/` before the response returns (and chains a best-effort
|
||||
/// AI overview run) — the per-entry fetch outcome is reported in the
|
||||
/// response's `source_fetch` field.
|
||||
#[serde(default)]
|
||||
source: Option<KnowledgeSource>,
|
||||
/// Optional tag keys to assign at creation time (same semantics as
|
||||
/// `UpdateBaseRequest.tags`).
|
||||
#[serde(default)]
|
||||
tags: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
async fn create_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<CreateBaseRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBaseInfo>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
// Detect whether the source is connector-backed (e.g. feishu) so we can
|
||||
// fire-and-forget a first sync after creation without blocking the response.
|
||||
let is_connector_source = req
|
||||
.source
|
||||
.as_ref()
|
||||
.is_some_and(|s| s.kind != "url" && !s.kind.is_empty());
|
||||
let mut info = state
|
||||
.service
|
||||
.create_base(&req.name, &req.description, req.root_path.as_deref(), req.source)
|
||||
.await?;
|
||||
// Persist tags (if provided) as a post-creation step — avoids changing the
|
||||
// 4-param `create_base` signature used by 50+ callers.
|
||||
if let Some(ref tag_keys) = req.tags {
|
||||
if !tag_keys.is_empty() {
|
||||
info = state.service.update_base(&info.id, None, None, Some(tag_keys.clone())).await?;
|
||||
}
|
||||
}
|
||||
// Connector-backed sources (feishu, etc.): trigger background sync so the
|
||||
// user does not have to manually invoke /sync after creation.
|
||||
if is_connector_source {
|
||||
let service = state.service.clone();
|
||||
let kb_id = info.id.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = service.sync_connector_source(&kb_id).await {
|
||||
tracing::warn!(kb_id, error = %e, "background connector sync after create failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(Json(ApiResponse::ok(info)))
|
||||
}
|
||||
|
||||
async fn get_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBaseInfo>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.get_base_info(&id).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct UpdateBaseRequest {
|
||||
name: Option<String>,
|
||||
description: Option<String>,
|
||||
#[serde(default)]
|
||||
tags: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
async fn update_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<UpdateBaseRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBaseInfo>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
state
|
||||
.service
|
||||
.update_base(&id, req.name.as_deref(), req.description.as_deref(), req.tags)
|
||||
.await?,
|
||||
)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DeleteBaseQuery {
|
||||
#[serde(default)]
|
||||
purge: bool,
|
||||
}
|
||||
|
||||
async fn delete_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
Query(query): Query<DeleteBaseQuery>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
state.service.delete_base(&id, query.purge).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
async fn list_files(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<KbFileEntry>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_files(&id).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ExportBaseRequest {
|
||||
/// Absolute destination path of the zip package.
|
||||
dest_path: String,
|
||||
}
|
||||
|
||||
async fn export_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<ExportBaseRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<ExportSummary>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
export::export_base(&state.service, &id, std::path::Path::new(&req.dest_path)).await?,
|
||||
)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ImportBaseRequest {
|
||||
/// Absolute path of a zip package created by the export endpoint.
|
||||
src_path: String,
|
||||
}
|
||||
|
||||
/// On success the service's managed-create path has already emitted
|
||||
/// `knowledge.base-created` (followed by `knowledge.base-updated` with the
|
||||
/// final file stats), so connected frontends refresh automatically. A
|
||||
/// best-effort AI overview run is then spawned in the background: it never
|
||||
/// overwrites a README carried by the package and only backfills the
|
||||
/// description when the package had none.
|
||||
async fn import_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<ImportBaseRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<ImportSummary>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let summary = export::import_base(&state.service, std::path::Path::new(&req.src_path)).await?;
|
||||
|
||||
let service = state.service.clone();
|
||||
let kb_id = summary.kb_id.clone();
|
||||
tokio::spawn(async move {
|
||||
// Best-effort: a missing completer (409) or an empty base (400) is
|
||||
// expected and must not surface anywhere. `None`: post-import
|
||||
// backfill is a background curation task → always the default model.
|
||||
if let Err(e) = service.generate_overview_opts(&kb_id, false, true, None).await {
|
||||
tracing::debug!(kb_id, error = %e, "post-import knowledge autogen skipped");
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Json(ApiResponse::ok(summary)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct AutogenRequest {
|
||||
/// Replace an existing `README.md`; default keeps it (the description is
|
||||
/// refreshed either way).
|
||||
#[serde(default)]
|
||||
overwrite_readme: bool,
|
||||
/// Explicit provider for the LLM call (the model picker). Must be sent
|
||||
/// together with `model` or not at all.
|
||||
#[serde(default)]
|
||||
provider_id: Option<String>,
|
||||
/// Explicit model for the LLM call. Must be sent together with
|
||||
/// `provider_id` or not at all.
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
/// Validate and assemble an optional explicit `(provider_id, model)` pick
|
||||
/// from a request: both fields must be present (non-blank) or both absent —
|
||||
/// a half-specified pick is a 400. Returns `Ok(None)` when neither is given
|
||||
/// (use the completer's default model).
|
||||
fn model_override(
|
||||
provider_id: Option<String>,
|
||||
model: Option<String>,
|
||||
) -> Result<Option<(String, String)>, AppError> {
|
||||
let provider_id = provider_id.map(|s| s.trim().to_owned()).filter(|s| !s.is_empty());
|
||||
let model = model.map(|s| s.trim().to_owned()).filter(|s| !s.is_empty());
|
||||
match (provider_id, model) {
|
||||
(Some(p), Some(m)) => Ok(Some((p, m))),
|
||||
(None, None) => Ok(None),
|
||||
_ => Err(AppError::BadRequest(
|
||||
"provider_id and model must be supplied together (or both omitted)".into(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// AI overview generation. Without a wired completer this returns 409 with
|
||||
/// an actionable message. Completion is broadcast as `knowledge.base-updated`.
|
||||
async fn autogen_base(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Option<Json<AutogenRequest>>,
|
||||
) -> Result<Json<ApiResponse<AutogenOutcome>>, AppError> {
|
||||
let req = body.map(|Json(r)| r).unwrap_or_default();
|
||||
let override_model = model_override(req.provider_id, req.model)?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
state.service.generate_overview(&id, req.overwrite_readme, override_model).await?,
|
||||
)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GenerateDescriptionRequest {
|
||||
/// Tentative base name from the create form; may be omitted/blank.
|
||||
#[serde(default)]
|
||||
name: String,
|
||||
/// Absolute path of an existing directory to sample.
|
||||
root_path: String,
|
||||
/// Explicit provider for the LLM call; pair with `model` or omit both.
|
||||
#[serde(default)]
|
||||
provider_id: Option<String>,
|
||||
/// Explicit model for the LLM call; pair with `provider_id` or omit both.
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct PolishDescriptionRequest {
|
||||
/// Tentative base name from the create form; may be omitted/blank.
|
||||
#[serde(default)]
|
||||
name: String,
|
||||
/// User-written draft description to rewrite.
|
||||
draft: String,
|
||||
/// Explicit provider for the LLM call; pair with `model` or omit both.
|
||||
#[serde(default)]
|
||||
provider_id: Option<String>,
|
||||
/// Explicit model for the LLM call; pair with `provider_id` or omit both.
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct DescriptionResponse {
|
||||
description: String,
|
||||
}
|
||||
|
||||
/// Stateless AI description generation for the create-base form: samples the
|
||||
/// given directory and returns a description only — no base row required,
|
||||
/// nothing persisted. 409 without a wired completer, 400 on an invalid path.
|
||||
async fn generate_description(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<GenerateDescriptionRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<DescriptionResponse>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let override_model = model_override(req.provider_id, req.model)?;
|
||||
let description = state
|
||||
.service
|
||||
.generate_description_for_path(&req.name, &req.root_path, override_model)
|
||||
.await?;
|
||||
Ok(Json(ApiResponse::ok(DescriptionResponse { description })))
|
||||
}
|
||||
|
||||
/// Stateless AI polish of a user-written draft description. 409 without a
|
||||
/// wired completer, 400 on an empty draft. Nothing persisted.
|
||||
async fn polish_description(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<PolishDescriptionRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<DescriptionResponse>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let override_model = model_override(req.provider_id, req.model)?;
|
||||
let description = state.service.polish_description(&req.name, &req.draft, override_model).await?;
|
||||
Ok(Json(ApiResponse::ok(DescriptionResponse { description })))
|
||||
}
|
||||
|
||||
/// Re-fetch every URL-source entry (overwriting old snapshots) and stamp
|
||||
/// `extra.source.last_fetched_at`.
|
||||
async fn refresh_source(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<RefreshSourceSummary>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.refresh_source(&id).await?)))
|
||||
}
|
||||
|
||||
/// Pull a connector-backed base's remote documents into `snapshots/` (Feishu
|
||||
/// wiki, …). Distinct from `refresh-source`, which is for URL sources.
|
||||
async fn sync_source(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<RefreshSourceSummary>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.sync_connector_source(&id).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SetSourceRequest {
|
||||
/// New source config, or `null` to detach the base's source.
|
||||
#[serde(default)]
|
||||
source: Option<KnowledgeSource>,
|
||||
}
|
||||
|
||||
/// Attach / replace / clear a base's source config (`extra.source`). Used to
|
||||
/// wire a connector (Feishu, …) onto an existing base. Does not fetch — the
|
||||
/// caller triggers sync afterward.
|
||||
async fn set_source(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<SetSourceRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBaseInfo>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(state.service.set_source(&id, req.source).await?)))
|
||||
}
|
||||
|
||||
async fn list_credentials(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
) -> Result<Json<ApiResponse<Vec<ConnectorCredentialSummary>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_credentials().await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateCredentialRequest {
|
||||
/// Connector discriminator: "feishu", …
|
||||
kind: String,
|
||||
name: String,
|
||||
/// Connector-specific secret payload (e.g. Feishu `{ app_id, app_secret }`).
|
||||
/// Probed against the remote before being encrypted at rest.
|
||||
payload: serde_json::Value,
|
||||
}
|
||||
|
||||
async fn create_credential(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<CreateCredentialRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<ConnectorCredentialSummary>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let summary = state.service.create_credential(&req.kind, &req.name, req.payload).await?;
|
||||
Ok(Json(ApiResponse::ok(summary)))
|
||||
}
|
||||
|
||||
async fn delete_credential(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
state.service.delete_credential(&id).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
/// Re-probe a stored credential against its remote (the UI "test connection"
|
||||
/// action). Returns the connector identity on success.
|
||||
async fn test_credential(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<ConnectorIdentity>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.test_credential(&id).await?)))
|
||||
}
|
||||
|
||||
// ── Tag CRUD routes ──────────────────────────────────────────────────────
|
||||
|
||||
async fn list_tags(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
) -> Result<Json<ApiResponse<Vec<KnowledgeTag>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_tags().await?)))
|
||||
}
|
||||
|
||||
async fn create_tag(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<CreateKnowledgeTagRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeTag>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
state.service.create_tag(&req.label, req.color).await?,
|
||||
)))
|
||||
}
|
||||
|
||||
async fn update_tag(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(key): Path<String>,
|
||||
body: Result<Json<UpdateKnowledgeTagRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeTag>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
state.service.update_tag(&key, req).await?,
|
||||
)))
|
||||
}
|
||||
|
||||
async fn delete_tag(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(key): Path<String>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
state.service.delete_tag(&key).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
// ── P4 inbox review + consumers ───────────────────────────────────────
|
||||
|
||||
async fn list_inbox(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<InboxEntry>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_inbox(&id).await?)))
|
||||
}
|
||||
|
||||
/// Total unreviewed staged proposals across all bases (sidebar red-dot signal).
|
||||
async fn pending_inbox_count(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
) -> Result<Json<ApiResponse<usize>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.count_pending_inbox().await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct InboxItemQuery {
|
||||
scope: String,
|
||||
path: String,
|
||||
}
|
||||
|
||||
async fn inbox_diff(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
Query(q): Query<InboxItemQuery>,
|
||||
) -> Result<Json<ApiResponse<InboxDiff>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.inbox_diff(&id, &q.scope, &q.path).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct InboxActionRequest {
|
||||
scope: String,
|
||||
path: String,
|
||||
}
|
||||
|
||||
async fn merge_inbox(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<InboxActionRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<InboxMergeResult>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(state.service.merge_inbox(&id, &req.scope, &req.path).await?)))
|
||||
}
|
||||
|
||||
async fn discard_inbox(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<InboxActionRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
state.service.discard_inbox(&id, &req.scope, &req.path).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct InboxBatchRequest {
|
||||
kb_id: String,
|
||||
scope: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct InboxBatchResult {
|
||||
count: usize,
|
||||
}
|
||||
|
||||
async fn merge_all_inbox(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<InboxBatchRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<InboxBatchResult>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let count = state.service.merge_all_inbox(&req.kb_id, req.scope.as_deref()).await?;
|
||||
Ok(Json(ApiResponse::ok(InboxBatchResult { count })))
|
||||
}
|
||||
|
||||
async fn discard_all_inbox(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<InboxBatchRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<InboxBatchResult>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
let count = state.service.discard_all_inbox(&req.kb_id, req.scope.as_deref()).await?;
|
||||
Ok(Json(ApiResponse::ok(InboxBatchResult { count })))
|
||||
}
|
||||
|
||||
async fn list_consumers(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<ConsumerInfo>>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.list_consumers(&id).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FilePathQuery {
|
||||
path: String,
|
||||
}
|
||||
|
||||
async fn read_file(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
Query(query): Query<FilePathQuery>,
|
||||
) -> Result<Json<ApiResponse<KbFileContent>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.read_file(&id, &query.path).await?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct WriteFileRequest {
|
||||
path: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
async fn write_file(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
body: Result<Json<WriteFileRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
state.service.write_file(&id, &req.path, &req.content).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
async fn delete_file(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path(id): Path<String>,
|
||||
Query(query): Query<FilePathQuery>,
|
||||
) -> Result<Json<ApiResponse<()>>, AppError> {
|
||||
state.service.delete_file(&id, &query.path).await?;
|
||||
Ok(Json(ApiResponse::ok(())))
|
||||
}
|
||||
|
||||
async fn get_binding(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path((kind, target_id)): Path<(String, String)>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBinding>>, AppError> {
|
||||
Ok(Json(ApiResponse::ok(state.service.get_binding(&kind, &target_id).await?)))
|
||||
}
|
||||
|
||||
async fn set_binding(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
Path((kind, target_id)): Path<(String, String)>,
|
||||
body: Result<Json<KnowledgeBinding>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<KnowledgeBinding>>, AppError> {
|
||||
let Json(binding) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
Ok(Json(ApiResponse::ok(
|
||||
state.service.set_binding(&kind, &target_id, binding).await?,
|
||||
)))
|
||||
}
|
||||
|
||||
// ─── Manual search (read-only, scoped) ───────────────────────────────────────
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SearchBasesRequest {
|
||||
kb_ids: Vec<String>,
|
||||
query: String,
|
||||
limit: Option<usize>,
|
||||
}
|
||||
|
||||
async fn search_bases(
|
||||
State(state): State<KnowledgeRouterState>,
|
||||
Extension(_user): Extension<CurrentUser>,
|
||||
body: Result<Json<SearchBasesRequest>, JsonRejection>,
|
||||
) -> Result<Json<ApiResponse<Vec<KnowledgeSearchHit>>>, AppError> {
|
||||
let Json(req) = body.map_err(|e| AppError::BadRequest(e.to_string()))?;
|
||||
// Scope: only search the caller-supplied kb_ids. Empty list → empty result
|
||||
// (search_bases already handles this, but we make it explicit).
|
||||
if req.kb_ids.is_empty() {
|
||||
return Ok(Json(ApiResponse::ok(Vec::new())));
|
||||
}
|
||||
let limit = req.limit.unwrap_or(20);
|
||||
let hits = state.service.search_bases(&req.kb_ids, &req.query, limit).await?;
|
||||
Ok(Json(ApiResponse::ok(hits)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use http_body_util::BodyExt;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::*;
|
||||
use crate::testutil::make_service;
|
||||
|
||||
fn test_app(data_dir: &std::path::Path) -> Router {
|
||||
let service = Arc::new(make_service(data_dir));
|
||||
// The auth middleware normally injects `CurrentUser`; tests attach
|
||||
// it directly as a request extension.
|
||||
knowledge_routes(KnowledgeRouterState::new(service)).layer(Extension(CurrentUser {
|
||||
id: "u1".into(),
|
||||
username: "u1".into(),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn json_body(resp: axum::response::Response) -> serde_json::Value {
|
||||
let bytes = resp.into_body().collect().await.unwrap().to_bytes();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
/// Wire-contract pin for workpath bindings (frontend Task 11): the
|
||||
/// target_id is sent as ONE percent-encoded segment
|
||||
/// (`encodeURIComponent(workpathKey)`, `/` → `%2F`). axum must match
|
||||
/// the single-segment route on the encoded path and hand the DECODED
|
||||
/// path to the handler; the service then canonicalizes spellings
|
||||
/// (trailing slash et al) onto one row.
|
||||
#[tokio::test]
|
||||
async fn binding_route_extracts_percent_encoded_workpath() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let app = test_app(dir.path());
|
||||
|
||||
// Write under a trailing-slash spelling…
|
||||
let set = Request::post("/api/knowledge/binding/workpath/%2FUsers%2Fme%2Fproj%2F")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
r#"{"enabled":true,"writeback":false,"kb_ids":["kb_x"]}"#,
|
||||
))
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(set).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
// …and read it back under the canonical spelling: same row.
|
||||
let get = Request::get("/api/knowledge/binding/workpath/%2FUsers%2Fme%2Fproj")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(get).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
let v = json_body(resp).await;
|
||||
assert_eq!(v["success"], true, "{v}");
|
||||
assert_eq!(v["data"]["enabled"], true, "{v}");
|
||||
assert_eq!(v["data"]["kb_ids"][0], "kb_x", "{v}");
|
||||
|
||||
// A never-bound workpath reads as the default (disabled) binding.
|
||||
let get = Request::get("/api/knowledge/binding/workpath/%2Felsewhere")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(get).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
let v = json_body(resp).await;
|
||||
assert_eq!(v["data"]["enabled"], false, "{v}");
|
||||
|
||||
// The default-workpath sentinel needs no encoding at all.
|
||||
let get = Request::get("/api/knowledge/binding/workpath/__default__")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.oneshot(get).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
/// An unknown binding kind stays a 400 — `workpath` is now accepted,
|
||||
/// arbitrary kinds are not.
|
||||
#[tokio::test]
|
||||
async fn binding_route_rejects_unknown_kind() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let app = test_app(dir.path());
|
||||
let get = Request::get("/api/knowledge/binding/nonsense/x")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app.oneshot(get).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
/// A half-specified model pick (only `provider_id`, or only `model`) is a
|
||||
/// 400 BadRequest — the validation runs before any completer/path work,
|
||||
/// so it fires even with no completer wired and an arbitrary root_path.
|
||||
#[tokio::test]
|
||||
async fn description_generate_rejects_half_specified_model() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let app = test_app(dir.path());
|
||||
|
||||
for body in [
|
||||
r#"{"name":"x","root_path":"/tmp/x","provider_id":"p1"}"#,
|
||||
r#"{"name":"x","root_path":"/tmp/x","model":"m1"}"#,
|
||||
] {
|
||||
let req = Request::post("/api/knowledge/description/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body))
|
||||
.unwrap();
|
||||
let resp = app.clone().oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "body={body}");
|
||||
let v = json_body(resp).await;
|
||||
assert!(
|
||||
v["error"].as_str().unwrap_or_default().contains("supplied together"),
|
||||
"{v}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The matching `model_override` unit contract: both or neither.
|
||||
#[test]
|
||||
fn model_override_requires_both_or_neither() {
|
||||
assert_eq!(model_override(None, None).unwrap(), None);
|
||||
assert_eq!(
|
||||
model_override(Some("p".into()), Some("m".into())).unwrap(),
|
||||
Some(("p".into(), "m".into()))
|
||||
);
|
||||
// Blank strings collapse to "absent" — so blank+blank is None, not an error.
|
||||
assert_eq!(model_override(Some(" ".into()), Some(" ".into())).unwrap(), None);
|
||||
assert!(model_override(Some("p".into()), None).is_err());
|
||||
assert!(model_override(None, Some("m".into())).is_err());
|
||||
assert!(model_override(Some("p".into()), Some(" ".into())).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_search_returns_scoped_hits() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let app = test_app(dir.path());
|
||||
|
||||
// 1. Create a knowledge base.
|
||||
let resp = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::post("/api/knowledge/bases")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(r#"{"name":"规范","description":""}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
let v = json_body(resp).await;
|
||||
let kb_id = v["data"]["id"].as_str().unwrap().to_owned();
|
||||
|
||||
// 2. Write a file with a keyword.
|
||||
let write_body = serde_json::json!({
|
||||
"path": "a.md",
|
||||
"content": "# 评审\n选中态用 primary-1"
|
||||
});
|
||||
let resp = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::put(format!("/api/knowledge/bases/{kb_id}/file"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(write_body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK, "write_file failed");
|
||||
|
||||
// 3. Search via the new route.
|
||||
let search_body = serde_json::json!({
|
||||
"kbIds": [kb_id],
|
||||
"query": "评审",
|
||||
"limit": 10
|
||||
});
|
||||
let resp = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::post("/api/knowledge/search")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(search_body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
let v = json_body(resp).await;
|
||||
let hits = v["data"].as_array().expect("data should be an array");
|
||||
assert!(
|
||||
hits.iter().any(|h| h["kb_id"].as_str() == Some(&kb_id)),
|
||||
"expected hit for kb_id={kb_id}, got {v}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_search_empty_kb_ids_returns_empty() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let app = test_app(dir.path());
|
||||
|
||||
let search_body = serde_json::json!({
|
||||
"kbIds": [],
|
||||
"query": "anything"
|
||||
});
|
||||
let resp = app
|
||||
.oneshot(
|
||||
Request::post("/api/knowledge/search")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(search_body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
let v = json_body(resp).await;
|
||||
let hits = v["data"].as_array().expect("data should be an array");
|
||||
assert!(hits.is_empty(), "empty kb_ids should return empty hits");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,802 @@
|
||||
//! URL knowledge source: SSRF-guarded fetching, HTML→Markdown conversion,
|
||||
//! and snapshot formatting for `{kb_root}/snapshots/{slug}.md` files.
|
||||
//!
|
||||
//! SSRF baseline: only http(s) URLs; the host is resolved BEFORE connecting
|
||||
//! and every resolved address must be public (loopback, private, link-local,
|
||||
//! CGNAT, unspecified, multicast and v4-mapped equivalents are rejected).
|
||||
//! The validated addresses are pinned onto the client (`resolve_to_addrs`)
|
||||
//! so the connection cannot re-resolve elsewhere, redirects are disabled in
|
||||
//! reqwest and followed manually (≤ [`MAX_REDIRECTS`] hops) with the full
|
||||
//! validation re-applied per hop.
|
||||
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use nomifun_common::AppError;
|
||||
use url::Url;
|
||||
|
||||
/// Base-root-relative directory holding URL snapshots.
|
||||
pub const SNAPSHOT_REL_DIR: &str = "snapshots";
|
||||
|
||||
/// Whole-request timeout per hop.
|
||||
pub const FETCH_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
/// Response bodies are truncated beyond this size.
|
||||
pub const FETCH_MAX_BYTES: usize = 5 * 1024 * 1024;
|
||||
/// Persisted snapshot bodies are truncated beyond this size (applies when no
|
||||
/// completer is available to condense an oversized page).
|
||||
pub const SNAPSHOT_MAX_BYTES: usize = 256 * 1024;
|
||||
/// Maximum manual redirect hops.
|
||||
pub const MAX_REDIRECTS: usize = 3;
|
||||
/// Slug length cap (ASCII chars).
|
||||
pub const SLUG_MAX_LEN: usize = 80;
|
||||
|
||||
/// A fetched page, converted to markdown.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FetchedPage {
|
||||
/// URL after redirects (the one the content actually came from).
|
||||
pub final_url: String,
|
||||
/// `<title>` of the page when it was HTML.
|
||||
pub title: Option<String>,
|
||||
pub markdown: String,
|
||||
/// True when the response body exceeded the size cap and was cut.
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// Page-fetching seam for knowledge URL sources (same late-wire pattern as
|
||||
/// [`crate::autogen::KnowledgeCompleter`]). The knowledge crate ships the
|
||||
/// trait plus its HTTP implementation ([`HttpFetcher`]); a heavier
|
||||
/// browser-rendering backend (`BrowserFetcher`) lives in `nomifun-ai-agent`
|
||||
/// and is late-wired via [`crate::service::KnowledgeService::with_url_fetcher`],
|
||||
/// so the knowledge crate never gains a browser-engine dependency (the P3
|
||||
/// anti-cycle decision ②).
|
||||
#[async_trait::async_trait]
|
||||
pub trait PageFetcher: Send + Sync {
|
||||
/// Fetch `raw_url` and return its markdown body (+ title / final URL /
|
||||
/// truncation flag). Same contract as the original `UrlFetcher::fetch_page`.
|
||||
async fn fetch_page(&self, raw_url: &str) -> Result<FetchedPage, AppError>;
|
||||
}
|
||||
|
||||
/// SSRF-guarded HTTP page fetcher (the first [`PageFetcher`] implementation;
|
||||
/// formerly `UrlFetcher`). Plain reqwest GET with HTML→markdown conversion —
|
||||
/// no JS rendering. `Default` uses the production limits; tests loosen them
|
||||
/// via the builder methods.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HttpFetcher {
|
||||
timeout: Duration,
|
||||
max_bytes: usize,
|
||||
allow_private: bool,
|
||||
}
|
||||
|
||||
/// Backward-compatible alias for the pre-trait name. External callers (e.g.
|
||||
/// `nomifun-gateway::tools_knowledge`) still refer to `UrlFetcher`; it is now
|
||||
/// the concrete HTTP implementation of [`PageFetcher`].
|
||||
pub type UrlFetcher = HttpFetcher;
|
||||
|
||||
impl Default for HttpFetcher {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
timeout: FETCH_TIMEOUT,
|
||||
max_bytes: FETCH_MAX_BYTES,
|
||||
allow_private: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl PageFetcher for HttpFetcher {
|
||||
async fn fetch_page(&self, raw_url: &str) -> Result<FetchedPage, AppError> {
|
||||
// Delegate to the inherent method so direct `HttpFetcher::fetch_page`
|
||||
// callers (no trait import needed) and `dyn PageFetcher` share one body.
|
||||
HttpFetcher::fetch_page(self, raw_url).await
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpFetcher {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn timeout(mut self, timeout: Duration) -> Self {
|
||||
self.timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn max_bytes(mut self, max_bytes: usize) -> Self {
|
||||
self.max_bytes = max_bytes;
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable the private/local address guard. ONLY for tests (mock HTTP
|
||||
/// servers bind to loopback).
|
||||
pub fn allow_private_for_tests(mut self) -> Self {
|
||||
self.allow_private = true;
|
||||
self
|
||||
}
|
||||
|
||||
/// Fetch `raw_url` and convert the response to markdown. Every hop is
|
||||
/// SSRF-validated; bodies larger than the cap are truncated, not failed.
|
||||
pub async fn fetch_page(&self, raw_url: &str) -> Result<FetchedPage, AppError> {
|
||||
let mut url = parse_fetch_url(raw_url)?;
|
||||
for _hop in 0..=MAX_REDIRECTS {
|
||||
let addrs = resolve_validated(&url, self.allow_private).await?;
|
||||
let response = self.send(&url, &addrs).await?;
|
||||
let status = response.status();
|
||||
|
||||
if status.is_redirection() {
|
||||
let location = response
|
||||
.headers()
|
||||
.get(reqwest::header::LOCATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or_else(|| AppError::BadGateway(format!("redirect without Location from {url}")))?;
|
||||
let next = url
|
||||
.join(location)
|
||||
.map_err(|e| AppError::BadGateway(format!("invalid redirect target {location}: {e}")))?;
|
||||
url = check_scheme(next)?;
|
||||
continue;
|
||||
}
|
||||
if !status.is_success() {
|
||||
return Err(AppError::BadGateway(format!("fetch failed: HTTP {status} for {url}")));
|
||||
}
|
||||
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
let (body, truncated) = self.read_capped(response).await?;
|
||||
let text = String::from_utf8_lossy(&body).into_owned();
|
||||
|
||||
let (title, markdown) = if looks_like_html(content_type.as_deref(), &text) {
|
||||
html_to_markdown(&text)
|
||||
} else {
|
||||
(None, text)
|
||||
};
|
||||
return Ok(FetchedPage {
|
||||
final_url: url.to_string(),
|
||||
title,
|
||||
markdown,
|
||||
truncated,
|
||||
});
|
||||
}
|
||||
Err(AppError::BadGateway(format!("too many redirects fetching {raw_url}")))
|
||||
}
|
||||
|
||||
async fn send(&self, url: &Url, addrs: &[SocketAddr]) -> Result<reqwest::Response, AppError> {
|
||||
// A fresh Client per hop is deliberate: `resolve_to_addrs` pins one
|
||||
// host's pre-validated addresses onto the client, and every redirect
|
||||
// hop may land on a different host needing its own pinning.
|
||||
let mut builder = nomifun_net::proxy::apply_detected_proxy(reqwest::Client::builder())
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.timeout(self.timeout);
|
||||
// Pin the pre-validated addresses so the actual connection cannot be
|
||||
// re-resolved to a different (private) host (DNS rebinding).
|
||||
if let Some(host) = url.host_str()
|
||||
&& !addrs.is_empty()
|
||||
{
|
||||
builder = builder.resolve_to_addrs(host, addrs);
|
||||
}
|
||||
let client = builder
|
||||
.build()
|
||||
.map_err(|e| AppError::Internal(format!("failed to build http client: {e}")))?;
|
||||
client
|
||||
.get(url.clone())
|
||||
.header(reqwest::header::USER_AGENT, "NomiFun-Knowledge/1.0")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
AppError::Timeout(format!("fetch timed out for {url}"))
|
||||
} else {
|
||||
AppError::BadGateway(format!("fetch failed for {url}: {e}"))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Drain the body up to `max_bytes`; longer bodies are truncated. A body
|
||||
/// of exactly `max_bytes` is kept whole and NOT flagged as truncated.
|
||||
async fn read_capped(&self, response: reqwest::Response) -> Result<(Vec<u8>, bool), AppError> {
|
||||
let mut body: Vec<u8> = Vec::new();
|
||||
let mut stream = response.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(c) => c,
|
||||
Err(e) if e.is_timeout() => return Err(AppError::Timeout(format!("fetch body timed out: {e}"))),
|
||||
Err(e) => return Err(AppError::BadGateway(format!("fetch body failed: {e}"))),
|
||||
};
|
||||
if body.len() + chunk.len() > self.max_bytes {
|
||||
let take = self.max_bytes - body.len();
|
||||
body.extend_from_slice(&chunk[..take]);
|
||||
return Ok((body, true));
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok((body, false))
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse + scheme-check a fetch URL (no DNS yet).
|
||||
fn parse_fetch_url(raw: &str) -> Result<Url, AppError> {
|
||||
let url = Url::parse(raw.trim()).map_err(|e| AppError::BadRequest(format!("invalid URL: {e}")))?;
|
||||
check_scheme(url)
|
||||
}
|
||||
|
||||
fn check_scheme(url: Url) -> Result<Url, AppError> {
|
||||
if !matches!(url.scheme(), "http" | "https") {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"only http(s) URLs are supported (got scheme: {})",
|
||||
url.scheme()
|
||||
)));
|
||||
}
|
||||
if url.host_str().is_none() {
|
||||
return Err(AppError::BadRequest("URL has no host".into()));
|
||||
}
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
/// Full pre-connect validation used by the fetcher and exposed for callers
|
||||
/// that want to vet a URL without fetching: scheme/host syntax plus a DNS
|
||||
/// resolution where EVERY resolved address must be public.
|
||||
pub async fn validate_fetch_url(raw: &str, allow_private: bool) -> Result<Url, AppError> {
|
||||
let url = parse_fetch_url(raw)?;
|
||||
resolve_validated(&url, allow_private).await?;
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
/// Resolve the URL host and reject private/local addresses. Returns the
|
||||
/// resolved socket addresses for connection pinning.
|
||||
async fn resolve_validated(url: &Url, allow_private: bool) -> Result<Vec<SocketAddr>, AppError> {
|
||||
let host = url
|
||||
.host_str()
|
||||
.ok_or_else(|| AppError::BadRequest("URL has no host".into()))?;
|
||||
let port = url.port_or_known_default().unwrap_or(443);
|
||||
|
||||
let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|e| AppError::BadGateway(format!("DNS resolution failed for {host}: {e}")))?
|
||||
.collect();
|
||||
if addrs.is_empty() {
|
||||
return Err(AppError::BadGateway(format!("DNS resolution returned no addresses for {host}")));
|
||||
}
|
||||
if !allow_private && let Some(bad) = addrs.iter().find(|a| forbidden_ip(&a.ip())) {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"URL host {host} resolves to a private or local address ({}); fetching it is blocked",
|
||||
bad.ip()
|
||||
)));
|
||||
}
|
||||
Ok(addrs)
|
||||
}
|
||||
|
||||
/// SSRF address policy: anything not unambiguously public is forbidden.
|
||||
fn forbidden_ip(ip: &IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(v4) => {
|
||||
let octets = v4.octets();
|
||||
v4.is_loopback()
|
||||
|| v4.is_private()
|
||||
|| v4.is_link_local()
|
||||
|| v4.is_unspecified()
|
||||
|| v4.is_broadcast()
|
||||
|| v4.is_multicast()
|
||||
|| v4.is_documentation()
|
||||
|| octets[0] == 0
|
||||
// CGNAT 100.64.0.0/10
|
||||
|| (octets[0] == 100 && (64..128).contains(&octets[1]))
|
||||
// IETF protocol assignments 192.0.0.0/24
|
||||
|| (octets[0] == 192 && octets[1] == 0 && octets[2] == 0)
|
||||
}
|
||||
IpAddr::V6(v6) => {
|
||||
let seg0 = v6.segments()[0];
|
||||
v6.is_loopback()
|
||||
|| v6.is_unspecified()
|
||||
|| v6.is_multicast()
|
||||
// Unique-local fc00::/7
|
||||
|| (seg0 & 0xfe00) == 0xfc00
|
||||
// Link-local fe80::/10
|
||||
|| (seg0 & 0xffc0) == 0xfe80
|
||||
// v4-mapped/compatible addresses inherit the v4 verdict.
|
||||
|| v6.to_ipv4_mapped().is_some_and(|v4| forbidden_ip(&IpAddr::V4(v4)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Decide whether a response body should go through HTML→MD conversion.
|
||||
/// The Content-Type header wins when it is conclusive; otherwise sniff the
|
||||
/// body prefix for an html document marker.
|
||||
fn looks_like_html(content_type: Option<&str>, body: &str) -> bool {
|
||||
if let Some(ct) = content_type {
|
||||
let ct = ct.to_ascii_lowercase();
|
||||
if ct.contains("html") {
|
||||
return true;
|
||||
}
|
||||
if ct.contains("markdown") || ct.contains("text/plain") || ct.contains("json") {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
let head: String = body.trim_start().chars().take(256).collect::<String>().to_ascii_lowercase();
|
||||
head.starts_with("<!doctype html") || head.starts_with("<html") || head.contains("<html")
|
||||
}
|
||||
|
||||
/// Convert HTML to markdown via `htmd`, falling back to `<title>` + stripped
|
||||
/// body text when conversion fails. Returns `(title, markdown)`.
|
||||
pub fn html_to_markdown(html: &str) -> (Option<String>, String) {
|
||||
let title = extract_html_title(html);
|
||||
let converter = htmd::HtmlToMarkdown::builder()
|
||||
.skip_tags(vec!["script", "style", "head", "iframe", "noscript"])
|
||||
.build();
|
||||
let markdown = match converter.convert(html) {
|
||||
Ok(md) if !md.trim().is_empty() => md,
|
||||
_ => {
|
||||
let mut text = strip_tags(html);
|
||||
if let Some(t) = &title {
|
||||
text = format!("# {t}\n\n{text}");
|
||||
}
|
||||
text
|
||||
}
|
||||
};
|
||||
(title, markdown)
|
||||
}
|
||||
|
||||
/// First `<title>…</title>` content, whitespace-collapsed.
|
||||
fn extract_html_title(html: &str) -> Option<String> {
|
||||
// ASCII-only lowercasing keeps byte offsets aligned with `html` (full
|
||||
// `to_lowercase` can change byte lengths, e.g. 'İ' → "i̇").
|
||||
let lower = html.to_ascii_lowercase();
|
||||
let open = lower.find("<title")?;
|
||||
let open_end = lower[open..].find('>').map(|i| open + i + 1)?;
|
||||
let close = lower[open_end..].find("</title").map(|i| open_end + i)?;
|
||||
let title = html.get(open_end..close)?;
|
||||
let title = title.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
(!title.is_empty()).then_some(title)
|
||||
}
|
||||
|
||||
/// Crude tag stripper used only as a conversion fallback: drops `<…>` spans
|
||||
/// and collapses blank-line runs.
|
||||
fn strip_tags(html: &str) -> String {
|
||||
let mut out = String::with_capacity(html.len() / 2);
|
||||
let mut in_tag = false;
|
||||
for c in html.chars() {
|
||||
match c {
|
||||
'<' => in_tag = true,
|
||||
'>' => {
|
||||
in_tag = false;
|
||||
out.push(' ');
|
||||
}
|
||||
c if !in_tag => out.push(c),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
let mut lines: Vec<&str> = Vec::new();
|
||||
let mut last_blank = true;
|
||||
for line in out.lines() {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
if !last_blank {
|
||||
lines.push("");
|
||||
}
|
||||
last_blank = true;
|
||||
} else {
|
||||
lines.push(trimmed);
|
||||
last_blank = false;
|
||||
}
|
||||
}
|
||||
lines.join("\n").trim().to_owned()
|
||||
}
|
||||
|
||||
/// Derive a snapshot file slug from the URL host+path: lowercase ASCII
|
||||
/// `[a-z0-9-]`, runs of other chars collapsed to single dashes, capped at
|
||||
/// [`SLUG_MAX_LEN`]. Never empty.
|
||||
pub fn slug_for_url(url: &Url) -> String {
|
||||
let raw = format!("{}{}", url.host_str().unwrap_or_default(), url.path());
|
||||
let mut slug = String::new();
|
||||
for c in raw.chars() {
|
||||
if slug.len() >= SLUG_MAX_LEN {
|
||||
break;
|
||||
}
|
||||
if c.is_ascii_alphanumeric() {
|
||||
slug.push(c.to_ascii_lowercase());
|
||||
} else if !slug.is_empty() && !slug.ends_with('-') {
|
||||
slug.push('-');
|
||||
}
|
||||
}
|
||||
let slug = slug.trim_matches('-').to_owned();
|
||||
if slug.is_empty() { "page".into() } else { slug }
|
||||
}
|
||||
|
||||
/// Assemble a snapshot document: YAML frontmatter (`source_url`,
|
||||
/// `fetched_at`, optional `title`) followed by the markdown body.
|
||||
pub fn snapshot_markdown(source_url: &str, fetched_at: &str, title: Option<&str>, body: &str) -> String {
|
||||
let mut out = String::from("---\n");
|
||||
out.push_str(&format!("source_url: {source_url}\n"));
|
||||
out.push_str(&format!("fetched_at: {fetched_at}\n"));
|
||||
if let Some(title) = title.map(str::trim).filter(|t| !t.is_empty()) {
|
||||
// Collapse runs of whitespace (incl. newlines/tabs): a multi-line
|
||||
// title would break out of its YAML frontmatter line.
|
||||
let title = title.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
out.push_str(&format!("title: \"{}\"\n", title.replace('"', "'")));
|
||||
}
|
||||
out.push_str("---\n\n");
|
||||
out.push_str(body.trim_end());
|
||||
out.push('\n');
|
||||
out
|
||||
}
|
||||
|
||||
/// Extract the `source_url` value from a snapshot's YAML frontmatter (the
|
||||
/// shape written by [`snapshot_markdown`]). Returns `None` for documents
|
||||
/// without a leading frontmatter block, or whose frontmatter has no
|
||||
/// `source_url` line — i.e. user-authored files that merely live in
|
||||
/// `snapshots/`. Only the frontmatter block is consulted; a `source_url:`
|
||||
/// line in the body never matches.
|
||||
pub fn snapshot_source_url(content: &str) -> Option<&str> {
|
||||
let rest = content.strip_prefix("---")?;
|
||||
let rest = rest.strip_prefix("\r\n").or_else(|| rest.strip_prefix('\n'))?;
|
||||
for line in rest.lines() {
|
||||
let trimmed = line.trim();
|
||||
if trimmed == "---" {
|
||||
return None; // frontmatter ended without the field
|
||||
}
|
||||
if let Some(value) = trimmed.strip_prefix("source_url:") {
|
||||
let value = value.trim();
|
||||
return (!value.is_empty()).then_some(value);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Truncate a string to at most `max_bytes`, never splitting a char.
|
||||
pub fn truncate_to_bytes(s: &str, max_bytes: usize) -> &str {
|
||||
if s.len() <= max_bytes {
|
||||
return s;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while end > 0 && !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
&s[..end]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
fn test_fetcher() -> UrlFetcher {
|
||||
UrlFetcher::new().allow_private_for_tests()
|
||||
}
|
||||
|
||||
// ── PageFetcher seam ─────────────────────────────────────────────
|
||||
|
||||
/// A non-HTTP [`PageFetcher`] (returns a canned page without touching the
|
||||
/// network) — proves the trait is object-safe and a custom backend can
|
||||
/// stand in for `HttpFetcher` behind `dyn PageFetcher` (the K2
|
||||
/// `BrowserFetcher` seam).
|
||||
struct CannedFetcher(FetchedPage);
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl PageFetcher for CannedFetcher {
|
||||
async fn fetch_page(&self, _raw_url: &str) -> Result<FetchedPage, AppError> {
|
||||
Ok(self.0.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_fetcher_is_usable_behind_dyn_page_fetcher() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/doc"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(
|
||||
"<html><head><title>渲染</title></head><body><h1>X</h1></body></html>",
|
||||
"text/html; charset=utf-8",
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
// Same code path, reached through the trait object rather than the
|
||||
// concrete type.
|
||||
let fetcher: std::sync::Arc<dyn PageFetcher> = std::sync::Arc::new(test_fetcher());
|
||||
let page = fetcher.fetch_page(&format!("{}/doc", server.uri())).await.unwrap();
|
||||
assert_eq!(page.title.as_deref(), Some("渲染"));
|
||||
assert!(page.markdown.contains("# X"), "got: {}", page.markdown);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn custom_page_fetcher_can_replace_http() {
|
||||
let fetcher: std::sync::Arc<dyn PageFetcher> = std::sync::Arc::new(CannedFetcher(FetchedPage {
|
||||
final_url: "https://spa.example.com/app".into(),
|
||||
title: Some("Rendered SPA".into()),
|
||||
markdown: "# Rendered\n\ncontent only a browser would see".into(),
|
||||
truncated: false,
|
||||
}));
|
||||
// The injected backend decides the result — no network involved.
|
||||
let page = fetcher.fetch_page("https://spa.example.com/app").await.unwrap();
|
||||
assert_eq!(page.title.as_deref(), Some("Rendered SPA"));
|
||||
assert!(page.markdown.contains("only a browser would see"));
|
||||
assert!(!page.truncated);
|
||||
}
|
||||
|
||||
// ── SSRF validation ──────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn validate_rejects_non_http_schemes() {
|
||||
for url in ["ftp://example.com/x", "file:///etc/passwd", "gopher://x", "javascript:alert(1)"] {
|
||||
let err = validate_fetch_url(url, false).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{url} → {err:?}");
|
||||
}
|
||||
assert!(validate_fetch_url("not a url", false).await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validate_rejects_loopback_private_and_linklocal() {
|
||||
for url in [
|
||||
"http://127.0.0.1/x",
|
||||
"http://127.8.8.8:9000/",
|
||||
"http://localhost/x",
|
||||
"http://10.0.0.5/",
|
||||
"http://172.16.1.1/",
|
||||
"http://192.168.1.1/admin",
|
||||
"http://169.254.169.254/latest/meta-data/",
|
||||
"http://100.64.0.1/",
|
||||
"http://0.0.0.0/",
|
||||
"http://[::1]/x",
|
||||
"http://[fe80::1]/",
|
||||
"http://[fc00::1]/",
|
||||
"http://[::ffff:127.0.0.1]/",
|
||||
] {
|
||||
let err = validate_fetch_url(url, false).await.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, AppError::BadRequest(_)),
|
||||
"{url} must be rejected, got {err:?}"
|
||||
);
|
||||
}
|
||||
// The test override admits loopback (mock servers).
|
||||
assert!(validate_fetch_url("http://127.0.0.1:1/x", true).await.is_ok());
|
||||
}
|
||||
|
||||
/// Obfuscated IPv4 literal notations (decimal, hex, octal): the url crate
|
||||
/// normalizes them all to dotted-quad form per the WHATWG URL spec, so
|
||||
/// the private-address guard must fire exactly as for `http://127.0.0.1/`.
|
||||
#[tokio::test]
|
||||
async fn validate_rejects_obfuscated_ipv4_literals() {
|
||||
for url in ["http://2130706433/", "http://0x7f000001/", "http://0177.0.0.1/"] {
|
||||
let err = validate_fetch_url(url, false).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{url} must be rejected, got {err:?}");
|
||||
}
|
||||
// Sanity-check the normalization assumption this test rests on.
|
||||
assert_eq!(Url::parse("http://2130706433/").unwrap().host_str(), Some("127.0.0.1"));
|
||||
assert_eq!(Url::parse("http://0x7f000001/").unwrap().host_str(), Some("127.0.0.1"));
|
||||
assert_eq!(Url::parse("http://0177.0.0.1/").unwrap().host_str(), Some("127.0.0.1"));
|
||||
}
|
||||
|
||||
/// Per-hop redirect validation, exercised at function level: `fetch_page`
|
||||
/// follows a redirect by joining the Location value onto the current URL
|
||||
/// (`Url::join`), re-checking the scheme (`check_scheme`) and re-resolving
|
||||
/// with the private-address guard (`resolve_validated`). An end-to-end
|
||||
/// wiremock test CANNOT cover the rejection: the mock server itself binds
|
||||
/// to loopback, so reaching hop 1 requires `allow_private` — which would
|
||||
/// also admit the private hop 2. This test drives the exact same functions
|
||||
/// on a redirect Location target instead.
|
||||
#[tokio::test]
|
||||
async fn redirect_hop_to_private_target_is_rejected() {
|
||||
let origin = Url::parse("https://public.example.com/start").unwrap();
|
||||
// Absolute Location to the cloud metadata endpoint (classic SSRF pivot).
|
||||
let next = origin.join("http://169.254.169.254/latest/meta-data/").unwrap();
|
||||
let next = check_scheme(next).unwrap();
|
||||
let err = resolve_validated(&next, false).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{err:?}");
|
||||
assert!(err.to_string().contains("private or local"), "{err}");
|
||||
|
||||
// A redirect downgrading to a non-http scheme dies at check_scheme.
|
||||
let bad = origin.join("ftp://internal/").unwrap();
|
||||
assert!(check_scheme(bad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forbidden_ip_policy() {
|
||||
let bad = ["127.0.0.1", "10.1.2.3", "172.31.0.1", "192.168.0.1", "169.254.0.1", "0.0.0.0", "100.100.0.1", "192.0.0.5", "224.0.0.1"];
|
||||
for ip in bad {
|
||||
assert!(forbidden_ip(&ip.parse().unwrap()), "{ip}");
|
||||
}
|
||||
let good = ["1.1.1.1", "8.8.8.8", "93.184.216.34", "100.128.0.1", "172.32.0.1"];
|
||||
for ip in good {
|
||||
assert!(!forbidden_ip(&ip.parse().unwrap()), "{ip}");
|
||||
}
|
||||
assert!(forbidden_ip(&"::1".parse().unwrap()));
|
||||
assert!(forbidden_ip(&"fe80::1".parse().unwrap()));
|
||||
assert!(forbidden_ip(&"fd12:3456::1".parse().unwrap()));
|
||||
assert!(forbidden_ip(&"::ffff:192.168.0.1".parse().unwrap()));
|
||||
assert!(!forbidden_ip(&"2606:4700:4700::1111".parse().unwrap()));
|
||||
}
|
||||
|
||||
// ── slug / frontmatter / conversion ──────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn slug_rules() {
|
||||
let u = |s: &str| Url::parse(s).unwrap();
|
||||
assert_eq!(slug_for_url(&u("https://docs.example.com/api/v2/Users")), "docs-example-com-api-v2-users");
|
||||
assert_eq!(slug_for_url(&u("https://example.com/")), "example-com");
|
||||
assert_eq!(slug_for_url(&u("https://example.com/a//b__c")), "example-com-a-b-c");
|
||||
let long = slug_for_url(&u(&format!("https://example.com/{}", "x".repeat(200))));
|
||||
assert!(long.len() <= SLUG_MAX_LEN, "{}", long.len());
|
||||
assert!(!long.ends_with('-'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontmatter_shape() {
|
||||
let md = snapshot_markdown(
|
||||
"https://example.com/docs",
|
||||
"2026-06-12T12:00:00Z",
|
||||
Some("My \"Docs\""),
|
||||
"# Title\n\nBody",
|
||||
);
|
||||
assert!(md.starts_with("---\nsource_url: https://example.com/docs\nfetched_at: 2026-06-12T12:00:00Z\n"), "got: {md}");
|
||||
assert!(md.contains("title: \"My 'Docs'\""), "got: {md}");
|
||||
assert!(md.contains("---\n\n# Title\n\nBody\n"), "got: {md}");
|
||||
// No title line when absent.
|
||||
let md = snapshot_markdown("https://e.com", "2026-01-01T00:00:00Z", None, "b");
|
||||
assert!(!md.contains("title:"), "got: {md}");
|
||||
// Newlines/tabs in a title collapse to single spaces — a multi-line
|
||||
// title must not break out of its frontmatter line.
|
||||
let md = snapshot_markdown("https://e.com", "2026-01-01T00:00:00Z", Some("Line one\nLine\ttwo"), "b");
|
||||
assert!(md.contains("title: \"Line one Line two\"\n"), "got: {md}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_source_url_reads_only_frontmatter() {
|
||||
// Round-trip with the writer.
|
||||
let md = snapshot_markdown("https://e.com/docs", "2026-01-01T00:00:00Z", Some("T"), "body");
|
||||
assert_eq!(snapshot_source_url(&md), Some("https://e.com/docs"));
|
||||
// User-authored files: no frontmatter at all, or frontmatter without
|
||||
// the field — a body-level `source_url:` line never counts.
|
||||
assert_eq!(snapshot_source_url("# notes\nsource_url: https://nope"), None);
|
||||
assert_eq!(snapshot_source_url("---\ntitle: x\n---\n\nsource_url: https://nope"), None);
|
||||
assert_eq!(snapshot_source_url("---\nsource_url:\n---\n"), None, "empty value is no value");
|
||||
assert_eq!(snapshot_source_url(""), None);
|
||||
// CRLF frontmatter still parses.
|
||||
assert_eq!(snapshot_source_url("---\r\nsource_url: https://e.com/x\r\n---\r\nbody"), Some("https://e.com/x"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn html_conversion_and_fallback() {
|
||||
let html = "<html><head><title>Guide Page</title><script>evil()</script></head>\
|
||||
<body><h1>Guide</h1><p>Hello <b>world</b></p></body></html>";
|
||||
let (title, md) = html_to_markdown(html);
|
||||
assert_eq!(title.as_deref(), Some("Guide Page"));
|
||||
assert!(md.contains("# Guide"), "got: {md}");
|
||||
assert!(md.contains("**world**"), "got: {md}");
|
||||
assert!(!md.contains("evil()"), "script content must be skipped: {md}");
|
||||
|
||||
// Tag stripping fallback keeps readable text.
|
||||
let text = strip_tags("<div><p>第一段</p>\n\n\n<p>第二段</p></div>");
|
||||
assert!(text.contains("第一段") && text.contains("第二段"), "got: {text}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_respects_char_boundaries() {
|
||||
let s = "知识库snapshot";
|
||||
let t = truncate_to_bytes(s, 4);
|
||||
assert_eq!(t, "知");
|
||||
assert_eq!(truncate_to_bytes("abc", 10), "abc");
|
||||
}
|
||||
|
||||
// ── fetching (mock HTTP) ─────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_converts_html_page() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/doc"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(
|
||||
"<html><head><title>接口文档</title></head><body><h1>API</h1><p>说明</p></body></html>",
|
||||
"text/html; charset=utf-8",
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let page = test_fetcher().fetch_page(&format!("{}/doc", server.uri())).await.unwrap();
|
||||
assert_eq!(page.title.as_deref(), Some("接口文档"));
|
||||
assert!(page.markdown.contains("# API"), "got: {}", page.markdown);
|
||||
assert!(page.markdown.contains("说明"));
|
||||
assert!(!page.truncated);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_passes_plaintext_through_and_truncates() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/big.md"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/plain")
|
||||
.set_body_string("x".repeat(4096)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let fetcher = test_fetcher().max_bytes(256);
|
||||
let page = fetcher.fetch_page(&format!("{}/big.md", server.uri())).await.unwrap();
|
||||
assert!(page.truncated);
|
||||
assert!(page.markdown.len() <= 256, "{}", page.markdown.len());
|
||||
assert!(page.title.is_none());
|
||||
}
|
||||
|
||||
/// A body of exactly `max_bytes` is kept whole and not flagged truncated.
|
||||
#[tokio::test]
|
||||
async fn fetch_body_exactly_at_cap_is_not_truncated() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/exact"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/plain")
|
||||
.set_body_string("x".repeat(256)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let fetcher = test_fetcher().max_bytes(256);
|
||||
let page = fetcher.fetch_page(&format!("{}/exact", server.uri())).await.unwrap();
|
||||
assert!(!page.truncated, "exactly-at-cap body must not be flagged");
|
||||
assert_eq!(page.markdown.len(), 256);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_follows_bounded_redirects() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/a"))
|
||||
.respond_with(ResponseTemplate::new(302).insert_header("location", "/b"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/b"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("landed"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
// Self-redirect loop.
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/loop"))
|
||||
.respond_with(ResponseTemplate::new(302).insert_header("location", "/loop"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let page = test_fetcher().fetch_page(&format!("{}/a", server.uri())).await.unwrap();
|
||||
assert!(page.markdown.contains("landed"));
|
||||
assert!(page.final_url.ends_with("/b"), "{}", page.final_url);
|
||||
|
||||
let err = test_fetcher().fetch_page(&format!("{}/loop", server.uri())).await.unwrap_err();
|
||||
assert!(err.to_string().contains("too many redirects"), "{err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_times_out_and_reports_http_errors() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/slow"))
|
||||
.respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(5)))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/missing"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let fetcher = test_fetcher().timeout(Duration::from_millis(200));
|
||||
let err = fetcher.fetch_page(&format!("{}/slow", server.uri())).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::Timeout(_)), "{err:?}");
|
||||
|
||||
let err = test_fetcher().fetch_page(&format!("{}/missing", server.uri())).await.unwrap_err();
|
||||
assert!(err.to_string().contains("404"), "{err}");
|
||||
}
|
||||
|
||||
/// Without the test override, fetching the loopback mock server must be
|
||||
/// rejected by the pre-connect guard (never reaches the socket).
|
||||
#[tokio::test]
|
||||
async fn fetch_blocks_private_targets_by_default() {
|
||||
let server = MockServer::start().await;
|
||||
let err = UrlFetcher::new().fetch_page(&format!("{}/doc", server.uri())).await.unwrap_err();
|
||||
assert!(matches!(err, AppError::BadRequest(_)), "{err:?}");
|
||||
assert!(err.to_string().contains("private or local"), "{err}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
//! Router state for the knowledge domain. Holds the `Arc`-wrapped service.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::service::KnowledgeService;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct KnowledgeRouterState {
|
||||
pub service: Arc<KnowledgeService>,
|
||||
}
|
||||
|
||||
impl KnowledgeRouterState {
|
||||
pub fn new(service: Arc<KnowledgeService>) -> Self {
|
||||
Self { service }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
//! Shared test fixtures for the knowledge crate: an in-memory
|
||||
//! `IKnowledgeRepository`, a no-op event broadcaster, and a service factory.
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use nomifun_common::TimestampMs;
|
||||
use nomifun_db::models::{CreateKnowledgeTagParams, KnowledgeBaseRow, KnowledgeBindingRow, KnowledgeTagRow, UpdateKnowledgeTagParams};
|
||||
use nomifun_db::{DbError, IKnowledgeRepository};
|
||||
|
||||
use crate::events::KnowledgeEventEmitter;
|
||||
use crate::service::KnowledgeService;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct MemRepo {
|
||||
pub bases: Mutex<Vec<KnowledgeBaseRow>>,
|
||||
/// Each binding row plus its ordered `kb_id` list (mirrors the
|
||||
/// `knowledge_binding_bases` junction).
|
||||
pub bindings: Mutex<Vec<(KnowledgeBindingRow, Vec<String>)>>,
|
||||
next_binding_id: AtomicI64,
|
||||
pub tags: Mutex<Vec<KnowledgeTagRow>>,
|
||||
}
|
||||
|
||||
/// Build a binding row with the `target_id` written to the column selected by
|
||||
/// `target_kind` (the in-memory analogue of the CHECK-constrained columns).
|
||||
fn binding_row(
|
||||
binding_id: i64,
|
||||
kind: &str,
|
||||
target_id: &str,
|
||||
enabled: bool,
|
||||
writeback: bool,
|
||||
writeback_mode: &str,
|
||||
writeback_eagerness: &str,
|
||||
channel_write_enabled: bool,
|
||||
updated_at: TimestampMs,
|
||||
) -> KnowledgeBindingRow {
|
||||
let mut row = KnowledgeBindingRow {
|
||||
binding_id,
|
||||
target_kind: kind.to_owned(),
|
||||
target_workpath: None,
|
||||
target_conv_id: None,
|
||||
target_term_id: None,
|
||||
target_companion_id: None,
|
||||
enabled,
|
||||
writeback,
|
||||
writeback_mode: writeback_mode.to_owned(),
|
||||
writeback_eagerness: writeback_eagerness.to_owned(),
|
||||
channel_write_enabled,
|
||||
updated_at,
|
||||
};
|
||||
// `target_conv_id` / `target_term_id` are integer FKs now; `target_workpath`
|
||||
// / `target_companion_id` stay string keys. Route the test id to the right column
|
||||
// with the right type.
|
||||
match kind {
|
||||
"workpath" => row.target_workpath = Some(target_id.to_owned()),
|
||||
"conversation" => row.target_conv_id = target_id.parse::<i64>().ok(),
|
||||
"terminal" => row.target_term_id = target_id.parse::<i64>().ok(),
|
||||
"companion" => row.target_companion_id = Some(target_id.to_owned()),
|
||||
_ => {}
|
||||
}
|
||||
row
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl IKnowledgeRepository for MemRepo {
|
||||
async fn insert_base(&self, row: &KnowledgeBaseRow) -> Result<(), DbError> {
|
||||
self.bases.lock().unwrap().push(row.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn update_base(&self, row: &KnowledgeBaseRow) -> Result<(), DbError> {
|
||||
let mut bases = self.bases.lock().unwrap();
|
||||
match bases.iter_mut().find(|r| r.id == row.id) {
|
||||
Some(r) => {
|
||||
*r = row.clone();
|
||||
Ok(())
|
||||
}
|
||||
None => Err(DbError::NotFound(row.id.clone())),
|
||||
}
|
||||
}
|
||||
async fn delete_base(&self, id: &str) -> Result<(), DbError> {
|
||||
let mut bases = self.bases.lock().unwrap();
|
||||
let before = bases.len();
|
||||
bases.retain(|r| r.id != id);
|
||||
if bases.len() == before {
|
||||
Err(DbError::NotFound(id.to_owned()))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
async fn get_base(&self, id: &str) -> Result<Option<KnowledgeBaseRow>, DbError> {
|
||||
Ok(self.bases.lock().unwrap().iter().find(|r| r.id == id).cloned())
|
||||
}
|
||||
async fn list_bases(&self) -> Result<Vec<KnowledgeBaseRow>, DbError> {
|
||||
Ok(self.bases.lock().unwrap().clone())
|
||||
}
|
||||
async fn get_binding(
|
||||
&self,
|
||||
kind: &str,
|
||||
id: &str,
|
||||
) -> Result<Option<(KnowledgeBindingRow, Vec<String>)>, DbError> {
|
||||
Ok(self
|
||||
.bindings
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|(b, _)| b.target_kind == kind && b.target_id().as_deref() == Some(id))
|
||||
.cloned())
|
||||
}
|
||||
async fn set_binding(
|
||||
&self,
|
||||
kind: &str,
|
||||
id: &str,
|
||||
kb_ids: &[String],
|
||||
enabled: bool,
|
||||
writeback: bool,
|
||||
writeback_mode: &str,
|
||||
writeback_eagerness: &str,
|
||||
channel_write_enabled: bool,
|
||||
updated_at: TimestampMs,
|
||||
) -> Result<i64, DbError> {
|
||||
let mut bindings = self.bindings.lock().unwrap();
|
||||
// Reuse the existing binding_id on upsert; allocate a fresh one
|
||||
// otherwise (the surrogate-key analogue of the real table).
|
||||
let binding_id = bindings
|
||||
.iter()
|
||||
.find(|(b, _)| b.target_kind == kind && b.target_id().as_deref() == Some(id))
|
||||
.map(|(b, _)| b.binding_id)
|
||||
.unwrap_or_else(|| self.next_binding_id.fetch_add(1, Ordering::SeqCst) + 1);
|
||||
bindings.retain(|(b, _)| !(b.target_kind == kind && b.target_id().as_deref() == Some(id)));
|
||||
let row = binding_row(binding_id, kind, id, enabled, writeback, writeback_mode, writeback_eagerness, channel_write_enabled, updated_at);
|
||||
bindings.push((row, kb_ids.to_vec()));
|
||||
Ok(binding_id)
|
||||
}
|
||||
async fn delete_binding(&self, kind: &str, id: &str) -> Result<(), DbError> {
|
||||
self.bindings
|
||||
.lock()
|
||||
.unwrap()
|
||||
.retain(|(b, _)| !(b.target_kind == kind && b.target_id().as_deref() == Some(id)));
|
||||
Ok(())
|
||||
}
|
||||
async fn list_bindings_using_kb(&self, kb_id: &str) -> Result<Vec<KnowledgeBindingRow>, DbError> {
|
||||
let mut rows: Vec<KnowledgeBindingRow> = self
|
||||
.bindings
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|(_, kb_ids)| kb_ids.iter().any(|k| k == kb_id))
|
||||
.map(|(b, _)| b.clone())
|
||||
.collect();
|
||||
rows.sort_by(|a, b| a.target_kind.cmp(&b.target_kind).then(a.binding_id.cmp(&b.binding_id)));
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_knowledge_tags(&self) -> Result<Vec<KnowledgeTagRow>, DbError> {
|
||||
let mut tags = self.tags.lock().unwrap().clone();
|
||||
tags.sort_by(|a, b| a.sort_order.cmp(&b.sort_order).then(a.key.cmp(&b.key)));
|
||||
Ok(tags)
|
||||
}
|
||||
|
||||
async fn create_knowledge_tag(&self, params: CreateKnowledgeTagParams) -> Result<(), DbError> {
|
||||
self.tags.lock().unwrap().push(KnowledgeTagRow {
|
||||
key: params.key,
|
||||
label: params.label,
|
||||
color: params.color,
|
||||
sort_order: params.sort_order,
|
||||
created_at: params.created_at,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn update_knowledge_tag(&self, key: &str, params: UpdateKnowledgeTagParams) -> Result<(), DbError> {
|
||||
let mut tags = self.tags.lock().unwrap();
|
||||
match tags.iter_mut().find(|t| t.key == key) {
|
||||
Some(t) => {
|
||||
if let Some(label) = params.label {
|
||||
t.label = label;
|
||||
}
|
||||
if let Some(color) = params.color {
|
||||
t.color = color;
|
||||
}
|
||||
if let Some(sort_order) = params.sort_order {
|
||||
t.sort_order = sort_order;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
None => Err(DbError::NotFound(format!("knowledge tag {key}"))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_knowledge_tag(&self, key: &str) -> Result<(), DbError> {
|
||||
let mut tags = self.tags.lock().unwrap();
|
||||
let before = tags.len();
|
||||
tags.retain(|t| t.key != key);
|
||||
if tags.len() == before {
|
||||
Err(DbError::NotFound(format!("knowledge tag {key}")))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct NoopBroadcaster;
|
||||
|
||||
impl nomifun_realtime::EventBroadcaster for NoopBroadcaster {
|
||||
fn broadcast(&self, _event: nomifun_api_types::WebSocketMessage<serde_json::Value>) {}
|
||||
}
|
||||
|
||||
pub(crate) fn make_service(data_dir: &Path) -> KnowledgeService {
|
||||
KnowledgeService::new(
|
||||
Arc::new(MemRepo::default()),
|
||||
data_dir,
|
||||
KnowledgeEventEmitter::new(Arc::new(NoopBroadcaster)),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
//! Workpath keys — the canonical identifier a `workpath`-kind knowledge
|
||||
//! binding is stored under (session-list unification spec §7).
|
||||
//!
|
||||
//! Knowledge mounts are physically per-workspace, so the binding belongs to
|
||||
//! the workspace path ("workpath"), not to the individual session. Every
|
||||
//! session whose workspace normalizes to the same key shares one binding
|
||||
//! row `('workpath', key)`; backend-managed temporary workspaces all map to
|
||||
//! the [`DEFAULT_WORKPATH_KEY`] sentinel.
|
||||
//!
|
||||
//! KEEP IN SYNC with the TypeScript twin
|
||||
//! `ui/src/renderer/pages/conversation/SessionList/utils/workpathKey.ts` —
|
||||
//! both sides must derive byte-identical keys or reads and writes land on
|
||||
//! different rows. The normalization rules (and the test set) are the same
|
||||
//! on both sides: trim → empty ⇒ sentinel → backslashes to forward slashes
|
||||
//! → a bare `/` root is kept → trailing slashes stripped. No case folding.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
/// Binding `target_kind` for workpath-level knowledge bindings.
|
||||
pub const WORKPATH_BINDING_KIND: &str = "workpath";
|
||||
|
||||
/// Sentinel `target_id` for the default workpath (backend-managed temporary
|
||||
/// workspaces). Same constant as the TS side's `DEFAULT_WORKPATH_KEY`.
|
||||
pub const DEFAULT_WORKPATH_KEY: &str = "__default__";
|
||||
|
||||
/// Normalize a workspace path into its binding key. Mirrors the TS
|
||||
/// `workpathKey()` exactly (see module docs): trim → empty ⇒
|
||||
/// [`DEFAULT_WORKPATH_KEY`] → `\` ⇒ `/` → root `/` kept → trailing slashes
|
||||
/// stripped.
|
||||
pub fn workpath_key(path: &str) -> String {
|
||||
let trimmed = path.trim();
|
||||
if trimmed.is_empty() {
|
||||
return DEFAULT_WORKPATH_KEY.to_owned();
|
||||
}
|
||||
let slashed = trimmed.replace('\\', "/");
|
||||
if slashed == "/" {
|
||||
return slashed;
|
||||
}
|
||||
slashed.trim_end_matches('/').to_owned()
|
||||
}
|
||||
|
||||
/// The workpath key a SESSION's workspace resolves to: a backend-managed
|
||||
/// (temporary) workspace — one under `managed_root` — is the default
|
||||
/// workpath ([`DEFAULT_WORKPATH_KEY`]); anything else keys by its
|
||||
/// normalized path.
|
||||
///
|
||||
/// The "is temporary" test matches the derivations the session DTOs already
|
||||
/// expose: conversations' `is_temporary_workspace`
|
||||
/// (`nomifun-conversation/src/convert.rs`, `workspace` under the backend
|
||||
/// data dir) and terminals' `is_default_workpath`
|
||||
/// (`nomifun-terminal/src/types.rs`, `cwd` under the backend work dir) —
|
||||
/// both guard against an empty root, since every path "starts with" an
|
||||
/// empty prefix.
|
||||
pub fn session_workpath_key(workspace: &Path, managed_root: &Path) -> String {
|
||||
if workspace.as_os_str().is_empty() {
|
||||
return DEFAULT_WORKPATH_KEY.to_owned();
|
||||
}
|
||||
if !managed_root.as_os_str().is_empty() && workspace.starts_with(managed_root) {
|
||||
return DEFAULT_WORKPATH_KEY.to_owned();
|
||||
}
|
||||
workpath_key(&workspace.to_string_lossy())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Same case set as the TS twin's `workpathKey.test.ts` — keep both in
|
||||
/// sync when normalization rules change.
|
||||
#[test]
|
||||
fn workpath_key_matches_ts_normalization() {
|
||||
// 去尾斜杠
|
||||
assert_eq!(workpath_key("/Users/a/proj/"), "/Users/a/proj");
|
||||
// 保留根路径
|
||||
assert_eq!(workpath_key("/"), "/");
|
||||
// 空白输入归 default
|
||||
assert_eq!(workpath_key(""), DEFAULT_WORKPATH_KEY);
|
||||
assert_eq!(workpath_key(" "), DEFAULT_WORKPATH_KEY);
|
||||
// 不做大小写折叠
|
||||
assert_eq!(workpath_key("/Users/A"), "/Users/A");
|
||||
// Windows 反斜杠归一为正斜杠并去尾
|
||||
assert_eq!(workpath_key("C:\\work\\proj\\"), "C:/work/proj");
|
||||
// 已规范化的 key 再归一化是幂等的(服务端二次归一化依赖这一点)。
|
||||
assert_eq!(workpath_key(DEFAULT_WORKPATH_KEY), DEFAULT_WORKPATH_KEY);
|
||||
assert_eq!(workpath_key("/Users/a/proj"), "/Users/a/proj");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_workpath_key_maps_temp_workspaces_to_default() {
|
||||
let root = Path::new("/data");
|
||||
// 临时工作区(位于 data_dir/work_dir 下)→ 哨兵。
|
||||
assert_eq!(
|
||||
session_workpath_key(Path::new("/data/conversations/gemini-temp-1"), root),
|
||||
DEFAULT_WORKPATH_KEY
|
||||
);
|
||||
// workspace == root 也算默认(starts_with 覆盖相等)。
|
||||
assert_eq!(session_workpath_key(root, root), DEFAULT_WORKPATH_KEY);
|
||||
// 空 workspace → 哨兵。
|
||||
assert_eq!(session_workpath_key(Path::new(""), root), DEFAULT_WORKPATH_KEY);
|
||||
// 用户自选目录 → 归一化 key(含去尾斜杠)。
|
||||
assert_eq!(session_workpath_key(Path::new("/Users/a/proj/"), root), "/Users/a/proj");
|
||||
// 空 root 不能把所有路径都吸成默认。
|
||||
assert_eq!(session_workpath_key(Path::new("/Users/a/proj"), Path::new("")), "/Users/a/proj");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user