Update: 将子项目从 submodule 转为完整内容

- 移除 GovAI, nomifun-tauri, 算力盒子 的 submodule 引用
- 添加所有子项目的完整源代码
- 保留原始 .git 为 .git.bak 备份
This commit is contained in:
freedak
2026-07-04 19:20:46 +08:00
parent 54d6465fa7
commit f7a720204a
3360 changed files with 802660 additions and 3 deletions
@@ -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(&not_zip, "definitely not a zip").unwrap();
let err = import_base(&service, &not_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");
}
}