diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 87890571d..cb6c2a72c 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -178,5 +178,8 @@ jobs: run: pnpm install --frozen-lockfile - name: Test run: pnpm test + # 测量台的判等不依赖任何包，借这个 job 现成的 node 跑，不另开一条流水线 + - name: Bench scoring tests + run: node --test ../scripts/bench/lib.test.mjs - name: Build (含类型检查) run: pnpm build diff --git a/README.md b/README.md index 9f0055035..e29fed9d9 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,7 @@ One Rust binary and one Postgres. Full-text search is embedded in the binary, ve | **Agent adjudication** | The agent adjudicates on its own, from common sense and the business documents. Low-confidence calls go to the human review queue, and the decisions people make there are recorded and used to tune the agent, closing the loop. | | **Reasoning and derivation** | Ontology axioms compile into rules: transitivity, symmetry, inverses and relation hierarchy derive new facts by forward chaining. Derivation is off by default, since a wrong axiom derives wrong facts. A derived fact is marked as such on the graph, carries validity and confidence like any other, and shows what it was derived from. When it contradicts an asserted fact, the asserted one stands. | | **Conflict detection** | Three kinds of conflict, three sets of choices. A new fact that clashes with an older one: close the old, keep both, or reject the new. Data that breaks an axiom (self-loop, asymmetry, transitive cycle, cardinality): retract the fact, relax the axiom, or accept both. The ontology itself is checked first, because violations of a self-contradictory ontology are noise. | -| **Ontology-driven querying** | Mount a database on a base (Postgres, MySQL and the engines that speak its protocol, Trino for Iceberg / Delta Lake / Hive, Databricks, Snowflake) and chat can query it alongside the documents. The agent proposes how its tables map onto the ontology, and you confirm. The method behind it, [Ontology2SQL](https://github.com/deeplethe/ontology2sql), is state of the art on BIRD Mini-Dev for SQLite and PostgreSQL ([submission](https://github.com/bird-bench/bird-bench.github.io/pull/218)). | +| **Ontology-driven querying** | Mount a database on a base (Postgres, MySQL and the engines that speak its protocol, ClickHouse, Trino for Iceberg / Delta Lake / Hive, Databricks, Snowflake) and chat can query it alongside the documents. The agent proposes how its tables map onto the ontology, and you confirm. The method behind it, [Ontology2SQL](https://github.com/deeplethe/ontology2sql), is state of the art on BIRD Mini-Dev for SQLite and PostgreSQL ([submission](https://github.com/bird-bench/bird-bench.github.io/pull/218)). | | **Multi-user and permissions** | Each knowledge base has its own members and roles: owner, admin, editor and viewer. Open bases are readable by everyone in the deployment, restricted ones only by invitation. The first account registered becomes the system administrator. | | **Decision ledger** | Confirming or rejecting a fact, merging or reverting an entity, rebuilding the graph: each leaves a record of who, when, and what the object looked like at the time. The ledger is append-only, and a record outlives its object, even the base it belonged to. | | **[Decision intelligence (in development)](#roadmap)** | Record a decision, replay both what was understood and the course it took, and reason over overlaid scenarios. | @@ -126,7 +126,7 @@ cd web && pnpm install && pnpm dev - [ ] **Business rules**: rules written by people over an entity's attribute facts, a threshold or a category set, that classify it as a derived fact with the rule and the premises as its explanation ([#277](https://github.com/deeplethe/utopia/issues/277)) - [ ] **Execution gate**: checking an agent's calls against ontology rules and symbolic logic - [ ] **MaxCompute**: mapping exploration and Ontology2SQL over Alibaba Cloud MaxCompute (Iceberg / Delta Lake via Trino, Databricks and Snowflake are in, awaiting a run against a real cluster) -- [ ] **More sources**: a ClickHouse driver; a Feishu connector +- [ ] **More sources**: a Feishu connector - [ ] **Agent memory over MCP**: episode writes, the retrieve endpoint, and the MCP server - [ ] **Enterprise**: OIDC SSO, backup and restore commands, benchmarks at 100k documents diff --git a/README.zh-CN.md b/README.zh-CN.md index d92fcadba..827247e72 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -127,7 +127,7 @@ cd web && pnpm install && pnpm dev - [ ] **业务规则**：由人写下、作用于实体属性事实的规则（阈值、类别集合），把实体归类为一条带前提的派生事实，规则与前提就是它的解释（[#277](https://github.com/deeplethe/utopia/issues/277)） - [ ] **执行校验层**：对 Agent 的调用进行本体规则与符号逻辑校验 - [ ] **问数与映射添加数据湖仓支持**：Iceberg / Delta Lake，以及 Databricks、Snowflake、MaxCompute 的映射探索与 Ontology2SQL 支持 -- [ ] **更多数据源**：ClickHouse 驱动，飞书连接器 +- [ ] **更多数据源**：飞书连接器 - [ ] **精确到时刻**：在年 / 月 / 日之外加一档 `instant` 精度，给那些本来就带时间戳的来源——现在连接器按 UTC 截到天，跨午夜的事件会差一天 - [ ] **MCP 上的 Agent 记忆**：补齐 episodes 写入、retrieve 端点与 MCP 服务器 - [ ] **企业化**：OIDC SSO、备份恢复命令、10 万文档级别的性能基准 diff --git a/crates/utopia-cli/src/main.rs b/crates/utopia-cli/src/main.rs index 3408e2576..2ec617196 100644 --- a/crates/utopia-cli/src/main.rs +++ b/crates/utopia-cli/src/main.rs @@ -82,7 +82,7 @@ struct ManifestDataDir { /// not a side effect of a code change. // 是迁移文件的**个数**，不是最大的编号（守卫 `schema_version_policy_compares_against_current` // 按个数比）：编号有空缺时两者不同——0071 由一个开放 PR 占着，0072 先落，个数是 71 -const CURRENT_SCHEMA_VERSION: u32 = 96; +const CURRENT_SCHEMA_VERSION: u32 = 97; fn main() -> anyhow::Result<()> { dotenvy::dotenv().ok(); diff --git a/crates/utopia-ingest/src/parsers.rs b/crates/utopia-ingest/src/parsers.rs index d715634e9..686b59b3c 100644 --- a/crates/utopia-ingest/src/parsers.rs +++ b/crates/utopia-ingest/src/parsers.rs @@ -485,6 +485,10 @@ pub(crate) fn docx_xml_to_text( /// PPTX: extract a:t text in the presentation's logical slide order. Tables under a:tbl /// become Markdown grids through the same renderer as DOCX and spreadsheet tables. +/// +/// 第一页是这份演示的开头，写在任何标题之前；第二页起每页一节（`## Slide N`）。封面上的 +/// 日期说的是整份演示，而一个日期管到哪看它所在的标题（0064）：封面也是一节的话，它的 +/// 日期只管封面，后面每一页的陈述都没有见证。Word 的标题块本来就在第一个标题之前，这里一样 pub fn pptx(bytes: &[u8]) -> anyhow::Result { let mut archive = zip::ZipArchive::new(Cursor::new(bytes.to_vec())).context("Failed to unzip pptx")?; @@ -515,7 +519,12 @@ pub fn pptx(bytes: &[u8]) -> anyhow::Result { for (num, name) in slides { let xml = pptx_part(&mut archive, &name)?; let text = pptx_xml_to_text(&xml)?; - if !text.trim().is_empty() { + if text.trim().is_empty() { + continue; + } + if num == 1 { + out.push_str(&format!("{text}\n")); + } else { out.push_str(&format!("\n## Slide {num}\n{text}\n")); } } @@ -1089,20 +1098,16 @@ fn pptx_xml_to_text(xml: &str) -> anyhow::Result { } "a:tr" if table_depth == 1 => row.clear(), "a:tc" if table_depth == 1 => { + // DrawingML 的合并属性属于 a:tc，后面的 a:tcPr 只描述样式。 cell = Some(TableCell { - span: 1, + span: attr(&e, "gridSpan") + .and_then(|v| v.parse().ok()) + .unwrap_or(1), + horizontal_merge: truthy(&e, "hMerge"), + vertical_merge: truthy(&e, "vMerge"), ..TableCell::default() }); } - "a:tcPr" if table_depth == 1 => { - if let Some(c) = cell.as_mut() { - if let Some(span) = attr(&e, "gridSpan").and_then(|v| v.parse().ok()) { - c.span = span; - } - c.horizontal_merge = truthy(&e, "hMerge"); - c.vertical_merge = truthy(&e, "vMerge"); - } - } "a:t" => { if cell.is_some() || table_depth == 0 { in_text = true; @@ -1119,15 +1124,6 @@ fn pptx_xml_to_text(xml: &str) -> anyhow::Result { "a:tblPr" if table_depth == 1 => { first_is_header = truthy(&e, "firstRow"); } - "a:tcPr" if table_depth == 1 => { - if let Some(c) = cell.as_mut() { - if let Some(span) = attr(&e, "gridSpan").and_then(|v| v.parse().ok()) { - c.span = span; - } - c.horizontal_merge = truthy(&e, "hMerge"); - c.vertical_merge = truthy(&e, "vMerge"); - } - } "a:br" => match cell.as_mut() { Some(c) => c.text.push(' '), None if table_depth == 0 => out.push('\n'), diff --git a/crates/utopia-ingest/tests/pptx_breaks.rs b/crates/utopia-ingest/tests/pptx_breaks.rs index 874cbb41f..b4918be50 100644 --- a/crates/utopia-ingest/tests/pptx_breaks.rs +++ b/crates/utopia-ingest/tests/pptx_breaks.rs @@ -56,5 +56,5 @@ fn runs_without_a_break_still_join() { let text = read( r#"Helloworld"#, ); - assert!(text.contains("## Slide 1\nHello\nworld\n"), "{text}"); + assert!(text.starts_with("Hello\nworld\n"), "{text}"); } diff --git a/crates/utopia-ingest/tests/pptx_order.rs b/crates/utopia-ingest/tests/pptx_order.rs index a4b19c1fe..f9f1747d9 100644 --- a/crates/utopia-ingest/tests/pptx_order.rs +++ b/crates/utopia-ingest/tests/pptx_order.rs @@ -46,10 +46,12 @@ fn parsed(parts: &Parts) -> String { fn logical_order_and_page_numbers_do_not_depend_on_zip_or_part_order() { let mut parts = deck(); let expected = parsed(&parts); + // 第一页是开头，没有标题；第二页起每页一节 assert!( - expected.contains("## Slide 1\nSECOND-CREATED & 中 "), + expected.starts_with("SECOND-CREATED & 中 "), "{expected}" ); + assert!(!expected.contains("## Slide 1"), "{expected}"); assert!(expected.contains("## Slide 2\nFIRST-CREATED"), "{expected}"); assert!(!expected.contains("UNLISTED")); assert!(!expected.contains("NOTES")); @@ -66,7 +68,7 @@ fn logical_order_and_page_numbers_do_not_depend_on_zip_or_part_order() { .replace("rIdA", "rIdB") .replace("temp", "rIdA"); let normal = parsed(&parts); - assert!(normal.contains("## Slide 1\nFIRST-CREATED")); + assert!(normal.starts_with("FIRST-CREATED"), "{normal}"); assert!(normal.contains("## Slide 2\nSECOND-CREATED")); } @@ -216,3 +218,13 @@ fn empty_presentations_and_manifest_free_legacy_packages_keep_their_behavior() { no_root.retain(|(n, _)| n != "_rels/.rels"); assert_eq!(parsed(&no_root), parsed(&deck())); } + +/// 封面写在任何标题之前：它说的日期管整份演示，不只管它自己那一页（0064 的范围看标题）。 +/// 第一页没有字的演示没有开头，其余各页照旧各是一节 +#[test] +fn the_first_slide_opens_the_deck_and_later_slides_are_sections() { + let text = parsed(&deck()); + let first_heading = text.find("## ").expect("later slides are sections"); + assert!(text[..first_heading].contains("SECOND-CREATED"), "{text}"); + assert!(text[first_heading..].starts_with("## Slide 2\n"), "{text}"); +} diff --git a/crates/utopia-ingest/tests/pptx_tables.rs b/crates/utopia-ingest/tests/pptx_tables.rs index a4c816629..d4f94ce8e 100644 --- a/crates/utopia-ingest/tests/pptx_tables.rs +++ b/crates/utopia-ingest/tests/pptx_tables.rs @@ -31,12 +31,12 @@ fn text_box(text: &str) -> String { } fn cell(text: &str) -> String { - cell_with_props(text, "") + cell_with_attrs(text, "") } -fn cell_with_props(text: &str, props: &str) -> String { +fn cell_with_attrs(text: &str, attrs: &str) -> String { format!( - r#"{text}"# + r#"{text}"# ) } @@ -48,8 +48,8 @@ fn cell_with_paragraphs(paragraphs: &[&str]) -> String { format!(r#"{paragraphs}"#) } -fn empty_cell(props: &str) -> String { - format!(r#""#) +fn empty_cell(attrs: &str) -> String { + format!(r#""#) } fn row(cells: &[String]) -> String { @@ -135,7 +135,7 @@ fn a_slide_table_that_is_not_a_grid_of_data_keeps_its_words() { let text = parse_slide(&table( first_row, &[row(&[ - cell_with_props("Notice for the board", r#"gridSpan="2""#), + cell_with_attrs("Notice for the board", r#"gridSpan="2""#), empty_cell(r#"hMerge="1""#), ])], )); @@ -186,43 +186,60 @@ fn a_title_and_text_around_a_slide_table_keep_their_lines() { #[test] fn grid_span_and_horizontal_merge_keep_the_grid_shape() { - let text = parse_slide(&table( + let body = table( true, &[ row(&[ - cell_with_props("Summary", r#"gridSpan="2""#), - cell_with_props("COVERED", r#"hMerge="1""#), + cell_with_attrs("Summary", r#"gridSpan="2""#), + cell_with_attrs("HORIZONTAL-COVERED", r#"hMerge="1""#), cell("Total"), ]), row(&[cell("Q1"), cell("1,200"), cell("1,350")]), ], - )); - - assert!( - text.contains("| | Summary | Total |\n| --- | --- | --- |\n| Q1 | 1,200 | 1,350 |"), - "{text}" ); - assert!(!text.contains("COVERED"), "{text}"); + + // Both XML event forms must leave the merge attributes read from a:tc intact. + for properties in ["", ""] { + let text = parse_slide(&body.replace("", properties)); + assert!(!text.contains("HORIZONTAL-COVERED"), "{properties}: {text}"); + assert!( + text.contains("| | Summary | Total |\n| --- | --- | --- |\n| Q1 | 1,200 | 1,350 |"), + "{properties}: {text}" + ); + } } #[test] fn a_vertical_merge_continuation_stays_an_empty_cell() { - let text = parse_slide(&table( + let body = table( true, &[ row(&[cell("Quarter"), cell("Revenue"), cell("YoY")]), - row(&[cell("Q1"), cell("1,200"), cell("+12%")]), - row(&[cell("Q2"), cell("1,350"), empty_cell(r#"vMerge="1""#)]), + row(&[ + cell("Q1"), + cell("1,200"), + cell_with_attrs("+12%", r#"rowSpan="2""#), + ]), + row(&[ + cell("Q2"), + cell("1,350"), + cell_with_attrs("VERTICAL-COVERED", r#"vMerge="1""#), + ]), row(&[cell("Q3"), cell("1,410"), cell("+7%")]), ], - )); - - assert!( - text.contains( - "| Quarter | Revenue | YoY |\n| --- | --- | --- |\n| Q1 | 1,200 | +12% |\n| Q2 | 1,350 | |\n| Q3 | 1,410 | +7% |" - ), - "{text}" ); + + // An already empty continuation would pass even if vMerge were ignored. + for properties in ["", ""] { + let text = parse_slide(&body.replace("", properties)); + assert!(!text.contains("VERTICAL-COVERED"), "{properties}: {text}"); + assert!( + text.contains( + "| Quarter | Revenue | YoY |\n| --- | --- | --- |\n| Q1 | 1,200 | +12% |\n| Q2 | 1,350 | |\n| Q3 | 1,410 | +7% |" + ), + "{properties}: {text}" + ); + } } #[test] diff --git a/crates/utopia-llm/src/lib.rs b/crates/utopia-llm/src/lib.rs index d14dc5ba0..ca2134fec 100644 --- a/crates/utopia-llm/src/lib.rs +++ b/crates/utopia-llm/src/lib.rs @@ -114,19 +114,19 @@ pub fn rate_limited(err: &anyhow::Error) -> Option<&RateLimited> { /// 实测一篇 32 块的文档在一小时里两类都撞上：502 四次、连接没送到两次。 pub fn transient(err: &anyhow::Error) -> Option<(&'static str, Option)> { if let Some(hit) = rate_limited(err) { - return Some(("端点限流", hit.retry_after)); + return Some(("LLM endpoint is rate limiting", hit.retry_after)); } if let Some(hit) = unavailable(err) { - return Some(("端点不可用", hit.retry_after)); + return Some(("LLM endpoint is unavailable", hit.retry_after)); } if err.chain().any(|e| e.is::()) { - return Some(("流断在半路", None)); + return Some(("LLM response stream was interrupted", None)); } let sending = err .chain() .find_map(|e| e.downcast_ref::()) .filter(|u| !u.0.is_timeout()); - sending.map(|_| ("请求没送到", None)) + sending.map(|_| ("LLM request could not be sent", None)) } /// 端点开口了又半路没了：流断在一句话中间，既没有 `[DONE]` 也没有 `finish_reason`。 @@ -2122,7 +2122,7 @@ data: {\"choices\":[{\"delta\":{\"content\":\"tail\"},\"finish_reason\":\"stop\" server.await.unwrap(); assert_eq!( crate::transient(&err).map(|(w, _)| w), - Some("流断在半路"), + Some("LLM response stream was interrupted"), "{err:#}" ); @@ -2422,7 +2422,7 @@ data: [DONE] assert!(crate::is_unreachable(&err), "{err:#}"); assert_eq!( crate::transient(&err).map(|(w, _)| w), - Some("请求没送到"), + Some("LLM request could not be sent"), "{err:#}" ); diff --git a/crates/utopia-server/src/api/alerts_routes.rs b/crates/utopia-server/src/api/alerts_routes.rs index 07c1acd39..5683a0529 100644 --- a/crates/utopia-server/src/api/alerts_routes.rs +++ b/crates/utopia-server/src/api/alerts_routes.rs @@ -16,7 +16,7 @@ use uuid::Uuid; use crate::auth::AuthUser; use crate::error::ApiResult; -use crate::state::AppState; +use crate::state::{AppEvent, AppState}; /// 一页几组。弹窗里放得下的量——再多就该翻页，而不是让人滚一屏。 const PAGE: i64 = 8; @@ -139,7 +139,12 @@ pub async fn stream( State(state): State, AuthUser(_user): AuthUser, ) -> ApiResult>>> { - let mut rx = state.events.subscribe(); + Ok(alert_event_stream(state.events.subscribe())) +} + +fn alert_event_stream( + mut rx: broadcast::Receiver, +) -> Sse>> { let stream = async_stream::stream! { loop { match rx.recv().await { @@ -147,10 +152,111 @@ pub async fn stream( yield Ok(Event::default().event("alert").data("{}")); } Ok(_) => continue, - Err(broadcast::error::RecvError::Lagged(_)) => continue, + // 唯一的告警可能已被其他类型的通知挤掉。结束响应，让 EventSource + // 重连后通过 onRecover 补刷，即使之后再没有新告警。 + Err(broadcast::error::RecvError::Lagged(_)) => return, Err(broadcast::error::RecvError::Closed) => return, } } }; - Ok(Sse::new(stream).keep_alive(KeepAlive::default())) + Sse::new(stream).keep_alive(KeepAlive::default()) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::response::IntoResponse; + use axum::{routing::get, Router}; + use futures_util::{FutureExt, StreamExt}; + use std::time::Duration; + + #[tokio::test] + async fn lagged_alert_stream_ends_when_non_alert_events_displace_its_only_alert() { + let (sender, _) = broadcast::channel(256); + let route_sender = sender.clone(); + let app = Router::new().route( + "/events", + get(move || { + let sender = route_sender.clone(); + async move { + let receiver = sender.subscribe(); + sender + .send(AppEvent { + kb_id: None, + kind: "alert", + document_id: None, + }) + .unwrap(); + for _ in 0..257 { + sender + .send(AppEvent { + kb_id: Some(Uuid::now_v7()), + kind: "document", + document_id: None, + }) + .unwrap(); + } + alert_event_stream(receiver) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let request = async { + let response = client + .get(format!("http://{address}/events")) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + response.text().await.unwrap() + }; + // 保留 sender：被过滤的剩余事件不能让漏掉告警的连接一直开着。 + let result = tokio::time::timeout(Duration::from_secs(2), request).await; + server.abort(); + assert!(result + .expect("a lagged alert SSE response must end") + .is_empty()); + } + + #[tokio::test] + async fn alert_stream_filters_non_alerts_and_never_exposes_alert_details() { + let (sender, receiver) = broadcast::channel(256); + let response = alert_event_stream(receiver).into_response(); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + let mut body = response.into_body().into_data_stream(); + assert!(body.next().now_or_never().is_none()); + sender + .send(AppEvent { + kb_id: Some(Uuid::now_v7()), + kind: "document", + document_id: None, + }) + .unwrap(); + assert!(body.next().now_or_never().is_none()); + sender + .send(AppEvent { + kb_id: Some(Uuid::now_v7()), + kind: "alert", + document_id: Some(Uuid::now_v7()), + }) + .unwrap(); + let frame = body.next().await.unwrap().unwrap(); + assert_eq!(frame.as_ref(), b"event: alert\ndata: {}\n\n"); + assert!(body.next().now_or_never().is_none()); + drop(sender); + assert!(body.next().await.is_none()); + } + + #[tokio::test] + async fn dropping_an_alert_response_releases_its_subscription() { + let (sender, receiver) = broadcast::channel(256); + let response = alert_event_stream(receiver).into_response(); + assert_eq!(sender.receiver_count(), 1); + drop(response); + assert_eq!(sender.receiver_count(), 0); + } } diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index c9a27acb4..8e12a9552 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -185,7 +185,7 @@ pub(super) fn tools_schema(can_write: bool, data_source_names: &[String]) -> ser }, "sql": { "type": "string", - "description": "One SELECT/WITH statement in the source's own SQL dialect (PostgreSQL, Trino, Databricks or Snowflake; the schema document names the engine)." + "description": "One SELECT/WITH statement in the source's own SQL dialect (PostgreSQL, MySQL, ClickHouse, Trino, Databricks or Snowflake; the schema document names the engine)." }, "purpose": { "type": "string", diff --git a/crates/utopia-server/src/api/datasource_routes.rs b/crates/utopia-server/src/api/datasource_routes.rs index b43f914ea..18d8900ec 100644 --- a/crates/utopia-server/src/api/datasource_routes.rs +++ b/crates/utopia-server/src/api/datasource_routes.rs @@ -58,7 +58,7 @@ pub async fn create( utopia_core::AppError::invalid( "unsupported_conn_scheme", format!( - "Connection string must start with one of: postgres://, mysql://, trino://, databricks://, snowflake:// (engines: {})", + "Connection string must start with one of: postgres://, mysql://, trino://, databricks://, snowflake://, clickhouse:// (engines: {})", crate::query_engine::ENGINES.join(", ") ), ) @@ -115,7 +115,7 @@ pub async fn test_conn( AppError::invalid( "unsupported_conn_scheme", format!( - "Connection string must start with one of: postgres://, mysql://, trino://, databricks://, snowflake:// (engines: {})", + "Connection string must start with one of: postgres://, mysql://, trino://, databricks://, snowflake://, clickhouse:// (engines: {})", crate::query_engine::ENGINES.join(", ") ), ) diff --git a/crates/utopia-server/src/api/events_routes.rs b/crates/utopia-server/src/api/events_routes.rs index ad2ffe9dc..aa99b28db 100644 --- a/crates/utopia-server/src/api/events_routes.rs +++ b/crates/utopia-server/src/api/events_routes.rs @@ -24,7 +24,13 @@ pub async fn kb_events( ) -> ApiResult>>> { utopia_store::access::require_kb(&state.pool, &user, kb_id, Role::Viewer).await?; - let mut rx = state.events.subscribe(); + Ok(kb_event_stream(state.events.subscribe(), kb_id)) +} + +fn kb_event_stream( + mut rx: broadcast::Receiver, + kb_id: Uuid, +) -> Sse>> { let stream = async_stream::stream! { loop { match rx.recv().await { @@ -35,13 +41,14 @@ pub async fn kb_events( .data(serde_json::to_string(&ev).unwrap_or_else(|_| "{}".into()))), None => continue, }, - // 消费落后被跳帧：无所谓，事件只是"该刷新了"的信号 - Err(broadcast::error::RecvError::Lagged(_)) => continue, + // 丢掉的可能是这个库或告警的唯一通知；余下事件可能全部被过滤。 + // 结束响应，让 EventSource 重连后通过 onRecover 补刷。 + Err(broadcast::error::RecvError::Lagged(_)) => return, Err(broadcast::error::RecvError::Closed) => return, } } }; - Ok(Sse::new(stream).keep_alive(KeepAlive::default())) + Sse::new(stream).keep_alive(KeepAlive::default()) } #[derive(Debug, PartialEq)] @@ -66,6 +73,125 @@ fn relay(ev: &AppEvent, kb_id: Uuid) -> Option { #[cfg(test)] mod tests { use super::*; + use axum::response::IntoResponse; + use axum::{routing::get, Router}; + use futures_util::{FutureExt, StreamExt}; + use std::time::Duration; + + #[tokio::test] + async fn lagged_kb_stream_ends_when_its_only_document_or_global_alert_is_lost() { + let (here, elsewhere) = (Uuid::now_v7(), Uuid::now_v7()); + for (kb_id, kind) in [(Some(here), "document"), (None, "alert")] { + let (sender, _) = broadcast::channel(256); + let route_sender = sender.clone(); + let app = Router::new().route( + "/events", + get(move || { + let sender = route_sender.clone(); + async move { + let receiver = sender.subscribe(); + sender + .send(AppEvent { + kb_id, + kind, + document_id: None, + }) + .unwrap(); + // 首次 poll 前确定性地挤掉唯一有效通知，剩下的全部会被过滤。 + for _ in 0..257 { + sender + .send(AppEvent { + kb_id: Some(elsewhere), + kind: "document", + document_id: None, + }) + .unwrap(); + } + kb_event_stream(receiver, here) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let request = async { + let response = client + .get(format!("http://{address}/events")) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + response.text().await.unwrap() + }; + // 保留 sender，确保真实 HTTP 响应结束来自 Lagged，而不是通道关闭。 + let result = tokio::time::timeout(Duration::from_secs(2), request).await; + server.abort(); + assert!(result + .expect("a lagged KB SSE response must end") + .is_empty()); + } + } + + #[tokio::test] + async fn kb_stream_keeps_filtering_and_alert_payloads_until_the_channel_closes() { + let (here, elsewhere) = (Uuid::now_v7(), Uuid::now_v7()); + let (sender, receiver) = broadcast::channel(256); + let response = kb_event_stream(receiver, here).into_response(); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + let mut body = response.into_body().into_data_stream(); + // 初次连接不凭空发刷新事件，也不结束响应。 + assert!(body.next().now_or_never().is_none()); + sender + .send(AppEvent { + kb_id: Some(elsewhere), + kind: "document", + document_id: None, + }) + .unwrap(); + assert!(body.next().now_or_never().is_none()); + + let document = AppEvent { + kb_id: Some(here), + kind: "document", + document_id: Some(Uuid::now_v7()), + }; + sender.send(document.clone()).unwrap(); + let frame = body.next().await.unwrap().unwrap(); + assert_eq!( + frame.as_ref(), + format!( + "event: document\ndata: {}\n\n", + serde_json::to_string(&document).unwrap() + ) + .as_bytes() + ); + + for kb_id in [None, Some(here), Some(elsewhere)] { + sender + .send(AppEvent { + kb_id, + kind: "alert", + document_id: Some(Uuid::now_v7()), + }) + .unwrap(); + let frame = body.next().await.unwrap().unwrap(); + assert_eq!(frame.as_ref(), b"event: alert\ndata: {}\n\n"); + } + assert!(body.next().now_or_never().is_none()); + drop(sender); + assert!(body.next().await.is_none()); + } + + #[tokio::test] + async fn dropping_a_kb_response_releases_its_subscription() { + let (sender, receiver) = broadcast::channel(256); + let response = kb_event_stream(receiver, Uuid::now_v7()).into_response(); + assert_eq!(sender.receiver_count(), 1); + drop(response); + assert_eq!(sender.receiver_count(), 0); + } #[test] fn a_kb_stream_carries_its_own_events_and_every_alert() { diff --git a/crates/utopia-server/src/api/ontology_routes.rs b/crates/utopia-server/src/api/ontology_routes.rs index 6e99e3824..bb85fa8b9 100644 --- a/crates/utopia-server/src/api/ontology_routes.rs +++ b/crates/utopia-server/src/api/ontology_routes.rs @@ -30,13 +30,12 @@ pub async fn get( Path(kb_id): Path, ) -> ApiResult> { require_kb(&state, &user, kb_id, Role::Viewer).await?; - let entity_types = utopia_store::ontology::entity_type_views(&state.pool, kb_id).await?; - let relation_types = utopia_store::ontology::relation_type_views(&state.pool, kb_id).await?; - let misses = utopia_store::ontology::list_misses(&state.pool, kb_id).await?; - // 已忽略的单列一路：抑制照旧（提案与自动扩本体只看上面那份）， - // 但让人看得见抑制掉了什么、现在涨到多少 - let dismissed_misses = - utopia_store::ontology::list_dismissed_misses(&state.pool, kb_id).await?; + let (entity_types, relation_types, misses, dismissed_misses) = tokio::try_join!( + utopia_store::ontology::entity_type_views(&state.pool, kb_id), + utopia_store::ontology::relation_type_views(&state.pool, kb_id), + utopia_store::ontology::list_misses(&state.pool, kb_id), + utopia_store::ontology::list_dismissed_misses(&state.pool, kb_id), + )?; Ok(Json(json!({ "entity_types": entity_types, "relation_types": relation_types, diff --git a/crates/utopia-server/src/api/question_routes.rs b/crates/utopia-server/src/api/question_routes.rs index ca4fcc9c5..2502e7244 100644 --- a/crates/utopia-server/src/api/question_routes.rs +++ b/crates/utopia-server/src/api/question_routes.rs @@ -188,7 +188,8 @@ pub async fn record_result( Ok(Json(json!({ "ok": true }))) } -/// 两个数（决定 5）：问题答对了几条；代理的提案人改过或拒掉的占几成 +/// 两个数（决定 5）：问题答对了几条；代理的提案人改过或拒掉的占几成。另带文档说了、 +/// 本体还放不下的类别词与短语形状各有多少，工作台靠它说本体还缺什么 pub async fn report( State(state): State, AuthUser(user): AuthUser, @@ -197,10 +198,12 @@ pub async fn report( utopia_store::access::require_kb(&state.pool, &user, kb_id, Role::Viewer).await?; let q = questions::report(&state.pool, kb_id).await?; let p = utopia_store::ontology::agent_proposal_report(&state.pool, kb_id).await?; + let uncovered = utopia_store::ontology::uncovered(&state.pool, kb_id).await?; let decided = p.adopted + p.rejected; let changed = p.adopted_edited + p.rejected; Ok(Json(json!({ "questions": q, + "uncovered": uncovered, "proposals": { "open": p.open, "adopted": p.adopted, diff --git a/crates/utopia-server/src/api/token_routes.rs b/crates/utopia-server/src/api/token_routes.rs index c56283e30..c287d91ee 100644 --- a/crates/utopia-server/src/api/token_routes.rs +++ b/crates/utopia-server/src/api/token_routes.rs @@ -5,8 +5,10 @@ use axum::extract::{Path, State}; use axum::Json; +use chrono::{DateTime, Duration, Utc}; use serde::Deserialize; use serde_json::json; +use utopia_core::{AppError, AppResult}; use uuid::Uuid; use crate::auth::AuthUser; @@ -35,14 +37,36 @@ fn default_days() -> i64 { 90 } +fn token_expires_at(now: DateTime, expires_in_days: i64) -> AppResult>> { + if expires_in_days < 0 { + return Err(AppError::invalid( + "bad_token_expiry", + "expires_in_days must be non-negative", + )); + } + if expires_in_days == 0 { + return Ok(None); + } + + let duration = Duration::try_days(expires_in_days).ok_or_else(|| { + AppError::invalid( + "bad_token_expiry", + "Token expiration duration is out of range", + ) + })?; + let expires_at = now.checked_add_signed(duration).ok_or_else(|| { + AppError::invalid("bad_token_expiry", "Token expiration date is out of range") + })?; + Ok(Some(expires_at)) +} + /// 发一枚。**明文只在这一次的响应里出现**,之后库里只有哈希。 pub async fn issue( State(state): State, AuthUser(user): AuthUser, Json(req): Json, ) -> ApiResult> { - let expires_at = (req.expires_in_days > 0) - .then(|| chrono::Utc::now() + chrono::Duration::days(req.expires_in_days)); + let expires_at = token_expires_at(Utc::now(), req.expires_in_days)?; let (view, plain) = utopia_store::tokens::issue( &state.pool, user.id, @@ -92,3 +116,7 @@ pub async fn revoke( .await; Ok(Json(json!({ "ok": true }))) } + +#[cfg(test)] +#[path = "token_routes_tests.rs"] +mod tests; diff --git a/crates/utopia-server/src/api/token_routes_tests.rs b/crates/utopia-server/src/api/token_routes_tests.rs new file mode 100644 index 000000000..5eec3c06c --- /dev/null +++ b/crates/utopia-server/src/api/token_routes_tests.rs @@ -0,0 +1,197 @@ +use super::{token_expires_at, IssueReq}; +use axum::body::{to_bytes, Body}; +use axum::http::{Request, StatusCode}; +use chrono::{DateTime, Duration, TimeZone, Utc}; +use serde_json::{json, Value}; +use std::sync::Arc; +use tower::ServiceExt; +use utopia_core::AppError; +use uuid::Uuid; + +#[test] +fn omitted_expiration_defaults_to_ninety_days() { + let request: IssueReq = serde_json::from_value(json!({ "name": "laptop" })).unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, 0).unwrap(); + + assert_eq!(request.expires_in_days, 90); + assert_eq!( + token_expires_at(now, request.expires_in_days).unwrap(), + Some(Utc.with_ymd_and_hms(2026, 4, 1, 12, 0, 0).unwrap()) + ); +} + +#[test] +fn zero_is_permanent_and_representable_positive_days_are_accepted() { + let now = Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, 0).unwrap(); + assert_eq!(token_expires_at(now, 0).unwrap(), None); + + for days in [1, 365, 1_000_000] { + let expires_at = token_expires_at(now, days).unwrap().unwrap(); + assert_eq!(expires_at.signed_duration_since(now), Duration::days(days)); + } + + // 上限来自日期能否表示，不是任意的业务天数限制。 + let last_valid_days = DateTime::::MAX_UTC + .signed_duration_since(now) + .num_days(); + assert!(token_expires_at(now, last_valid_days).unwrap().is_some()); + assert!(token_expires_at(now, last_valid_days + 1).is_err()); +} + +#[test] +fn negative_days_and_both_overflow_paths_are_validation_errors() { + let now = Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, 0).unwrap(); + let date_overflow = Duration::try_days(100_000_000).unwrap(); + assert!(now.checked_add_signed(date_overflow).is_none()); + assert!(Duration::try_days(i64::MAX).is_none()); + + for (days, expected_message) in [ + (-1, "expires_in_days must be non-negative"), + (i64::MIN, "expires_in_days must be non-negative"), + (100_000_000, "Token expiration date is out of range"), + (i64::MAX, "Token expiration duration is out of range"), + ] { + let error = token_expires_at(now, days).unwrap_err(); + match error { + AppError::Invalid { + code, + message, + detail, + } => { + assert_eq!(code, "bad_token_expiry", "days={days}"); + assert_eq!(message, expected_message, "days={days}"); + assert!(detail.is_none()); + } + other => panic!("expected validation error for days={days}, got {other:?}"), + } + } +} + +#[tokio::test] +async fn issue_route_validates_expiration_before_writing_tokens_or_success_audits( +) -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = sqlx::PgPool::connect(&url).await?; + utopia_store::db::migrate(&pool).await?; + let org_id = Uuid::now_v7(); + let user_id = Uuid::now_v7(); + let directory = tempfile::tempdir()?; + let config = utopia_core::config::AppConfig { + data_dir: directory.path().to_string_lossy().into_owned(), + ..Default::default() + }; + let search = Arc::new(utopia_search::SearchIndex::open( + &directory.path().join("search"), + )?); + let state = crate::state::AppState::new(pool.clone(), &config, search, "test-only".into()); + let session_token = crate::auth::issue_token(&state, user_id)?; + let app = crate::api::router(state, &config); + + let result = async { + sqlx::query("INSERT INTO organizations (id, name) VALUES ($1, 'token-expiry-test')") + .bind(org_id) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO users (id, org_id, email, password_hash, display_name) + VALUES ($1, $2, $3, 'unused', 'Token Expiry Test')", + ) + .bind(user_id) + .bind(org_id) + .bind(format!("{user_id}@token-expiry.test")) + .execute(&pool) + .await?; + + let mut issued_count = 0_i64; + for (requested_days, accepted_days) in [ + (None, Some(90)), + (Some(0), Some(0)), + (Some(1), Some(1)), + (Some(365), Some(365)), + (Some(-1), None), + (Some(i64::MIN), None), + (Some(100_000_000), None), + (Some(i64::MAX), None), + ] { + let mut payload = json!({ "name": "laptop" }); + if let Some(days) = requested_days { + payload["expires_in_days"] = json!(days); + } + let request = Request::builder() + .method("POST") + .uri("/api/v1/me/tokens") + .header("authorization", format!("Bearer {session_token}")) + .header("content-type", "application/json") + .body(Body::from(payload.to_string()))?; + let before = Utc::now(); + let response = app.clone().oneshot(request).await?; + let after = Utc::now(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 65536).await?; + let body: Value = serde_json::from_slice(&bytes)?; + + if let Some(days) = accepted_days { + anyhow::ensure!(status == StatusCode::OK, "{payload}: {status} {body}"); + anyhow::ensure!(body["token"] + .as_str() + .is_some_and(|token| token.starts_with("utp_pat_"))); + anyhow::ensure!(body["info"]["scope"] == "read"); + anyhow::ensure!(body["info"]["kb_ids"].is_null()); + let expires_at: Option> = + serde_json::from_value(body["info"]["expires_at"].clone())?; + if days == 0 { + anyhow::ensure!(expires_at.is_none()); + } else { + let expires_at = + expires_at.ok_or_else(|| anyhow::anyhow!("missing expiration"))?; + // PostgreSQL 保留微秒，比较时使用同一精度。 + let earliest = (before + Duration::days(days)).timestamp_micros(); + let latest = (after + Duration::days(days)).timestamp_micros(); + anyhow::ensure!((earliest..=latest).contains(&expires_at.timestamp_micros())); + } + issued_count += 1; + } else { + anyhow::ensure!( + status == StatusCode::UNPROCESSABLE_ENTITY, + "{payload}: {status} {body}" + ); + anyhow::ensure!(body["code"] == "bad_token_expiry"); + anyhow::ensure!(body["error"] + .as_str() + .is_some_and(|message| !message.is_empty())); + anyhow::ensure!(body.get("token").is_none()); + } + + let token_count: i64 = + sqlx::query_scalar("SELECT count(*) FROM personal_tokens WHERE user_id = $1") + .bind(user_id) + .fetch_one(&pool) + .await?; + let audit_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM audit_events WHERE actor_id = $1 AND action = 'token.issued'", + ) + .bind(user_id) + .fetch_one(&pool) + .await?; + anyhow::ensure!( + token_count == issued_count, + "unexpected token count for {payload}" + ); + anyhow::ensure!( + audit_count == issued_count, + "unexpected success audit count for {payload}" + ); + } + anyhow::Ok(()) + } + .await; + + // 审计只追加，保留在专用测试库中；随机用户 ID 隔离每次运行的计数。 + sqlx::query("DELETE FROM organizations WHERE id = $1") + .bind(org_id) + .execute(&pool) + .await?; + result +} diff --git a/crates/utopia-server/src/extraction.rs b/crates/utopia-server/src/extraction.rs index fcd234983..7894eb6de 100644 --- a/crates/utopia-server/src/extraction.rs +++ b/crates/utopia-server/src/extraction.rs @@ -74,7 +74,7 @@ pub(crate) async fn chat_retrying_rate_limits_at( return Err(err); }; if attempt == RATE_LIMIT_TRIES { - return Err(err.context(format!("{what}退避 {RATE_LIMIT_TRIES} 次仍未通过"))); + return Err(err.context(format!("{what} after {RATE_LIMIT_TRIES} attempts"))); } let delay = jitter(retry_after.unwrap_or(backoff).min(RATE_LIMIT_CAP)); tracing::warn!( @@ -159,14 +159,15 @@ pub(crate) fn incomplete_reason(unextracted: &[(i32, String)], attempted: usize) .collect(); let more = unextracted.len().saturating_sub(sample.len()); let tail = if more > 0 { - format!("；另有 {more} 个") + format!("; {more} more") } else { String::new() }; Some(format!( - "本轮 {attempted} 个分块里 {} 个没能抽取：{}{tail}", + "{} of {attempted} {} could not be extracted in this attempt: {}{tail}", unextracted.len(), - sample.join("；") + if attempted == 1 { "chunk" } else { "chunks" }, + sample.join("; ") )) } @@ -419,7 +420,7 @@ mod tests { #[test] fn a_document_with_a_skipped_chunk_is_not_complete() { assert_eq!(incomplete_reason(&[], 23), None, "全抽完才算完成"); - let one = [(7, "调用失败：timeout".to_string())]; + let one = [(7, "request failed: timeout".to_string())]; let msg = incomplete_reason(&one, 23).expect("有块没抽成就不该算完成"); assert!(msg.contains("23"), "分母要说出来：{msg}"); assert!(msg.contains("#7"), "得指得出是哪一块：{msg}"); @@ -429,11 +430,11 @@ mod tests { /// 但**剩下多少必须说**——否则读的人会以为只坏了三块。 #[test] fn many_failures_are_summarised_without_hiding_the_count() { - let many: Vec<(i32, String)> = (1..=20).map(|i| (i, "调用失败".into())).collect(); + let many: Vec<(i32, String)> = (1..=20).map(|i| (i, "request failed".into())).collect(); let msg = incomplete_reason(&many, 60).unwrap(); assert!(msg.contains("20"), "总数要在：{msg}"); - assert!(msg.contains("另有 17 个"), "省略掉的数量要说出来：{msg}"); - assert_eq!(msg.matches("调用失败").count(), 3, "只举三个"); + assert!(msg.contains("; 17 more"), "省略掉的数量要说出来：{msg}"); + assert_eq!(msg.matches("request failed").count(), 3, "只举三个"); } /// 重试时 `chunks_for_extraction` 只取还没抽的块，所以分母是**本轮**的数， @@ -441,7 +442,10 @@ mod tests { #[test] fn the_denominator_is_this_rounds_chunks_not_the_document() { let msg = incomplete_reason(&[(2, "x".into())], 3).unwrap(); - assert!(msg.starts_with("本轮 3 个分块"), "{msg}"); + assert!( + msg.starts_with("1 of 3 chunks could not be extracted in this attempt:"), + "{msg}" + ); } #[tokio::test] diff --git a/crates/utopia-server/src/extraction_open.rs b/crates/utopia-server/src/extraction_open.rs index 6b6227518..adf877e2e 100644 --- a/crates/utopia-server/src/extraction_open.rs +++ b/crates/utopia-server/src/extraction_open.rs @@ -95,11 +95,11 @@ async fn call_and_parse_one_chunk( kb_id, document_id, reason::CHUNK_UNEXTRACTED, - "调用失败，这一块没有进图", - Some(&format!("#{}：{e}", chunk.seq)), + "request failed; this chunk did not enter the graph", + Some(&format!("#{}: {e}", chunk.seq)), ) .await; - unextracted.push((chunk.seq, format!("调用失败：{e}"))); + unextracted.push((chunk.seq, format!("request failed: {e}"))); return Ok(None); } }; @@ -117,13 +117,13 @@ async fn call_and_parse_one_chunk( // 别的解析失败是「回复不合结构」。两种都写成同一句话就分不开了 let (why, detail) = if cut_by_ceiling { ( - "回复撞上 token 上限被截断，剩下的解析不了，这一块没有进图", - format!("#{}：hit the token ceiling；{e}", chunk.seq), + "response hit the token ceiling and could not be parsed; this chunk did not enter the graph", + format!("#{}: hit the token ceiling; {e}", chunk.seq), ) } else { ( - "回复解析不了，这一块没有进图", - format!("#{}：{e}", chunk.seq), + "response could not be parsed; this chunk did not enter the graph", + format!("#{}: {e}", chunk.seq), ) }; drop_signal( @@ -135,7 +135,7 @@ async fn call_and_parse_one_chunk( Some(&detail), ) .await; - unextracted.push((chunk.seq, format!("回复解析失败：{e}"))); + unextracted.push((chunk.seq, format!("response parsing failed: {e}"))); return Ok(None); } }; @@ -696,19 +696,47 @@ pub(crate) fn dated_at(doc: &Document) -> Option> } /// 一段字在块里的**字符**偏移（起、止）。字符不是字节：界面和 SQL 的 `substr` 都按字符数， -/// 中文一个字三个字节，按字节存的偏移到界面上就错位。找不到原样的就 `None`—— -/// 偏移只能由服务端从原文算出来，模型报的数字不算数 +/// 中文一个字三个字节，按字节存的偏移到界面上就错位。找不到的就 `None`—— +/// 偏移只能由服务端从原文算出来，模型报的数字不算数。 +/// +/// **一段空白对一段空白**，别的逐字比。PDF 的文字层在行尾换行，「December⏎31, 2024」 +/// 模型抄回来是「December 31, 2024」；逐字比的话这句引文和它的日期都对不上，陈述落了、 +/// 时间却丢了（一份英文董事会纪要四个日期全没了）。换行还是空格不是原文说了什么， +/// 偏移照原文里的位置算 fn locate(hay: &str, needle: &str) -> Option<(i32, i32)> { - let needle = needle.trim(); + let (needle, _) = squashed(needle.trim()); if needle.is_empty() { return None; } - let byte = hay.find(needle)?; - let start = hay[..byte].chars().count(); - let end = start + needle.chars().count(); + let (hay, at) = squashed(hay); + let found = hay + .windows(needle.len()) + .position(|w| w == needle.as_slice())?; + let start = at[found]; + let end = at[found + needle.len() - 1] + 1; Some((start as i32, end as i32)) } +/// 连着的空白压成一个空格。第二项是压过之后每个字符在原文里的字符位置 +fn squashed(text: &str) -> (Vec, Vec) { + let (mut out, mut at) = (Vec::new(), Vec::new()); + let mut in_space = false; + for (i, c) in text.chars().enumerate() { + if c.is_whitespace() { + if !in_space { + out.push(' '); + at.push(i); + } + in_space = true; + } else { + out.push(c); + at.push(i); + in_space = false; + } + } + (out, at) +} + /// 时间词在块里的字符起点。**必须在这条陈述自己的那句引文里**：模型会把一个时间词挂到 /// 好几条陈述上（FDA 语料实测「week 4」挂到了「不应由过敏患者服用」上，六条带时间的陈述错了三条）， /// 整块里搜得到不等于这句说了它。引文里有、但引文本身没在块里定位到的，起点退回整块里的第一处 @@ -730,7 +758,16 @@ fn locate_time(chunk: &str, quote: Option<(&str, Option<(i32, i32)>)>, words: &s let (q, span) = quote?; if let Some((inner, _)) = locate(q, words) { return match span { - Some((start, _)) => Some(start + inner), + // 在原文的这一句里找：引文和原文的空白未必一样长（换行加缩进对一个空格）， + // 按引文里的位置加上去会偏 + Some((start, end)) => { + let sentence: String = chunk + .chars() + .skip(start.max(0) as usize) + .take((end - start).max(0) as usize) + .collect(); + Some(start + locate(&sentence, words).map_or(inner, |(i, _)| i)) + } None => locate(chunk, words).map(|(s, _)| s), }; } @@ -1192,6 +1229,38 @@ mod tests { assert_eq!(locate(text, ""), None); } + /// PDF 的文字层在行尾换行，模型抄回来是空格：一段空白对一段空白，偏移照原文算 + #[test] + fn a_line_break_in_the_passage_matches_a_space_in_the_copy() { + let text = "1. Succession. Zhang Wei will step down as Chief Technology Officer on December\n 31, 2024. Li Na follows."; + let quote = "Zhang Wei will step down as Chief Technology Officer on December 31, 2024."; + let span = locate(text, quote).expect("the sentence is in the passage"); + let found: String = text + .chars() + .skip(span.0 as usize) + .take((span.1 - span.0) as usize) + .collect(); + assert_eq!( + found, + "Zhang Wei will step down as Chief Technology Officer on December\n 31, 2024." + ); + // 时间词在原文那一句里的位置，不是在引文里的位置：换行加缩进比一个空格长 + let start = locate_time(text, Some((quote, Some(span))), "December 31, 2024").unwrap(); + assert_eq!( + text.chars() + .skip(start as usize) + .take(8) + .collect::(), + "December" + ); + // 反过来也一样：原文是一个空格，抄回来多了空白 + assert_eq!(locate("a b", "a \t b"), Some((0, 3))); + // 空白之外逐字比：断在词中间的、换了字的都不算 + assert_eq!(locate("Septem\nber 1, 2024", "September 1, 2024"), None); + assert_eq!(locate(text, "December 30, 2024"), None); + assert_eq!(locate(" ", " "), None); + } + #[test] fn a_time_is_found_inside_its_own_sentence_first() { let text = "In 2019 the plant opened. In 2019 it closed again."; diff --git a/crates/utopia-server/src/ingest_sources.rs b/crates/utopia-server/src/ingest_sources.rs index b0cedc1de..62f81c69d 100644 --- a/crates/utopia-server/src/ingest_sources.rs +++ b/crates/utopia-server/src/ingest_sources.rs @@ -1165,7 +1165,7 @@ async fn sync_notion(state: &AppState, source: &Source) -> anyhow::Result anyhow::Result>() + .join("; ") + ); + } Ok(stats) } diff --git a/crates/utopia-server/src/notion.rs b/crates/utopia-server/src/notion.rs index 59062a4f7..e53f1390b 100644 --- a/crates/utopia-server/src/notion.rs +++ b/crates/utopia-server/src/notion.rs @@ -42,11 +42,12 @@ const NOTION_VERSION: &str = "2026-03-11"; /// 可以有几万页。 const MAX_PAGES_PER_SYNC: usize = 500; -/// 每页最多取多少个 block。再深的页面截断，比让一次同步卡在一页上好。 +/// 每页所有层级合计最多取多少个 block，包含无文字容器。 +/// 再深的页面截断，比让一次同步卡在一页上好。 const MAX_BLOCKS_PER_PAGE: usize = 500; /// 两次请求之间至少隔这么久。Notion 说的是「平均每秒三次」，取 350 毫秒留一点余量： -/// 500 页的上限下一次同步最坏三分钟出头，比同步失败便宜得多。 +/// 总耗时还取决于每页的嵌套层级与分页数。 const MIN_INTERVAL: Duration = Duration::from_millis(350); /// 一次请求撞上 429 最多等几回。`Retry-After` 通常是个位数秒，连等几回还在限流 /// 就不是节奏问题了，该把错误交出去。 @@ -82,11 +83,17 @@ fn client(token: &str) -> anyhow::Result { struct Paced { http: reqwest::Client, last: Option, + // 私有地址入口让 HTTP 测试走真实取页和限流逻辑，不增加用户配置。 + api_root: String, } impl Paced { fn new(http: reqwest::Client) -> Self { - Self { http, last: None } + Self { + http, + last: None, + api_root: "https://api.notion.com/v1".into(), + } } /// 发一次请求并解析 JSON。`what` 进日志和错误文案（"search" / "blocks"）。 @@ -160,9 +167,22 @@ fn retry_after(header: Option<&str>) -> Duration { /// /// **只搜页面，不搜 data source。** 后者是表格的容器，它自己没有正文； /// 表格里的每一行是一个页面，会在同一次搜索里出现。 -pub async fn fetch(token: &str, query: Option<&str>) -> anyhow::Result<(Vec, bool)> { +/// +/// 第三项是正文没读出来的页，每页一句原因。它们不在第一项里，见 `fetch_pages`。 +pub async fn fetch( + token: &str, + query: Option<&str>, +) -> anyhow::Result<(Vec, bool, Vec)> { let mut http = Paced::new(client(token)?); + fetch_pages(&mut http, query).await +} + +async fn fetch_pages( + http: &mut Paced, + query: Option<&str>, +) -> anyhow::Result<(Vec, bool, Vec)> { let mut out = Vec::new(); + let mut unread: Vec = Vec::new(); let mut cursor: Option = None; let mut truncated = false; @@ -178,11 +198,8 @@ pub async fn fetch(token: &str, query: Option<&str>) -> anyhow::Result<(Vec) -> anyhow::Result<(Vec text, + Err(error) => { + tracing::warn!(%id, error = %format!("{error:#}"), "notion page body could not be read, leaving the stored page as it is"); + unread.push(format!("notion page {id}: {error:#}")); + continue; + } + }; out.push(NotionPage { external_key: format!("notion://{id}"), @@ -219,7 +244,7 @@ pub async fn fetch(token: &str, query: Option<&str>) -> anyhow::Result<(Vec String { /// 取一页的正文，逐层展开 block。 async fn page_text(http: &mut Paced, page_id: &str) -> anyhow::Result { let mut out = String::new(); - let mut n = 0usize; + let mut block_count = 0; + append_children(http, page_id, &mut out, &mut block_count).await?; + Ok(out) +} + +/// 每个父块独立分页，先写父块，再读完子树，最后继续兄弟块。 +/// 输出和预算由整页共享；递归深度也受同一块数上限约束。 +async fn append_children( + http: &mut Paced, + parent_id: &str, + out: &mut String, + block_count: &mut usize, +) -> anyhow::Result<()> { let mut cursor: Option = None; - loop { - let mut url = format!("https://api.notion.com/v1/blocks/{page_id}/children?page_size=100"); + while *block_count < MAX_BLOCKS_PER_PAGE { + let mut url = format!( + "{}/blocks/{parent_id}/children?page_size=100", + http.api_root + ); if let Some(c) = &cursor { url.push_str(&format!("&start_cursor={c}")); } let v = http.send("blocks", |c| c.get(&url)).await?; for b in v["results"].as_array().unwrap_or(&vec![]) { - if n >= MAX_BLOCKS_PER_PAGE { - return Ok(out); + if *block_count >= MAX_BLOCKS_PER_PAGE { + return Ok(()); } - n += 1; + *block_count += 1; if let Some(line) = render_block(b) { out.push_str(&line); out.push('\n'); } + // 不能因父块没有 rich_text 而跳过 column 等容器。 + // 独立页面由 search 单独摄入，数据库也不是父页正文的一部分。 + let independent_content = + matches!(b["type"].as_str(), Some("child_page" | "child_database")); + if *block_count < MAX_BLOCKS_PER_PAGE + && b["has_children"].as_bool() == Some(true) + && !independent_content + { + let id = b["id"] + .as_str() + .context("notion blocks: child block is missing id")?; + if b["type"].as_str() == Some("table") { + let mut rows = String::new(); + Box::pin(append_children(http, id, &mut rows, block_count)).await?; + out.push_str(&render_table(&b["table"], &rows)?); + } else { + Box::pin(append_children(http, id, out, block_count)).await?; + } + } } if v["has_more"].as_bool() != Some(true) { break; @@ -279,14 +338,14 @@ async fn page_text(http: &mut Paced, page_id: &str) -> anyhow::Result { break; } } - Ok(out) + Ok(()) } /// 把一个 block 渲染成一行文本。 /// /// **认不出的类型返回它的纯文本而不是丢掉。** Notion 的 block 类型一直在加， -/// 硬编码一张白名单意味着新类型静默消失；而所有带文字的 block 都把文字放在 -/// `{type}.rich_text` 下，这个形状很稳。 +/// 硬编码一张白名单意味着新类型静默消失；普通文字在 `{type}.rich_text`， +/// 表格行则按 `table_row.cells` 保留每个单元格。 fn render_block(b: &serde_json::Value) -> Option { let t = b["type"].as_str()?; let inner = &b[t]; @@ -307,6 +366,20 @@ fn render_block(b: &serde_json::Value) -> Option { let lang = inner["language"].as_str().unwrap_or(""); format!("```{lang}\n{text}\n```") } + "table_row" => { + let cells = inner["cells"].as_array()?; + let cells: Vec = cells + .iter() + .map(|cell| { + rich_text(cell) + .replace('\\', "\\\\") + .replace('|', "\\|") + .replace("\r\n", "\n") + .replace(['\r', '\n'], "
") + }) + .collect(); + format!("| {} |", cells.join(" | ")) + } // 分割线与图片没有 rich_text，但它们在正文里也没有信息量 "divider" | "image" | "video" | "file" => return None, // child_page 的标题在 `title` 而不是 rich_text @@ -316,6 +389,34 @@ fn render_block(b: &serde_json::Value) -> Option { }) } +/// 无列头的表用空表头占位，不能把第一条数据误标为列头。 +fn render_table(table: &serde_json::Value, rows: &str) -> anyhow::Result { + if rows.is_empty() { + return Ok(String::new()); + } + let width = table["table_width"] + .as_u64() + .context("notion table: missing column count")?; + let width = usize::try_from(width)?; + anyhow::ensure!(width > 0, "notion table: column count must be positive"); + let has_header = table["has_column_header"].as_bool() == Some(true); + let separator = format!("|{}\n", " --- |".repeat(width)); + let mut out = String::from("\n"); + if !has_header { + out.push_str(&format!("|{}\n", " |".repeat(width))); + out.push_str(&separator); + } + for (index, row) in rows.lines().enumerate() { + out.push_str(row); + out.push('\n'); + if index == 0 && has_header { + out.push_str(&separator); + } + } + out.push('\n'); + Ok(out) +} + /// rich_text 数组拼成纯文本。 fn rich_text(v: &serde_json::Value) -> String { v.as_array() @@ -341,6 +442,14 @@ fn slug(title: &str) -> String { } } +#[cfg(test)] +#[path = "notion_tests.rs"] +mod traversal_tests; + +#[cfg(test)] +#[path = "notion_resync_tests.rs"] +mod resync_tests; + #[cfg(test)] mod tests { use super::*; @@ -424,7 +533,7 @@ mod tests { eprintln!("跳过：未设 UTOPIA_NOTION_TEST_TOKEN"); return Ok(()); }; - let (pages, _) = fetch(&token, None).await?; + let (pages, _, _) = fetch(&token, None).await?; assert!( !pages.is_empty(), "一页都没有——integration 可能没有被分享任何页面" diff --git a/crates/utopia-server/src/notion_resync_tests.rs b/crates/utopia-server/src/notion_resync_tests.rs new file mode 100644 index 000000000..7fe9051ec --- /dev/null +++ b/crates/utopia-server/src/notion_resync_tests.rs @@ -0,0 +1,218 @@ +//! A page imported before nested traversal must gain its missing body on resync, +//! even when Notion's edit timestamp has not changed. A later failed read must +//! leave that complete body and its processing history unchanged. +use super::{fetch_pages, Paced}; +use crate::ingest_sources::{ingest_item, IngestAction}; +use chrono::{DateTime, Utc}; +use serde_json::json; +use std::sync::Arc; +use utopia_store::documents; +use uuid::Uuid; +use wiremock::{matchers::method, matchers::path, Mock, MockServer, ResponseTemplate}; + +#[tokio::test] +async fn resync_adds_nested_body_and_preserves_it_when_a_later_read_fails() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = sqlx::PgPool::connect(&url).await?; + let (org, workspace, kb, source) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + let result = async { + sqlx::query("INSERT INTO organizations(id,name) VALUES($1,'notion-resync-test')") + .bind(org) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO workspaces(id,org_id,name) VALUES($1,$2,'notion-resync-test')") + .bind(workspace) + .bind(org) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO knowledge_bases(id,workspace_id,name) VALUES($1,$2,'notion-resync-test')", + ) + .bind(kb) + .bind(workspace) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO sources(id,kb_id,kind,name) VALUES($1,$2,'notion','fixture')") + .bind(source) + .bind(kb) + .execute(&pool) + .await?; + let dir = tempfile::tempdir()?; + let config = utopia_core::config::AppConfig { + data_dir: dir.path().to_string_lossy().into_owned(), + ..Default::default() + }; + let search = Arc::new(utopia_search::SearchIndex::open( + &dir.path().join("search"), + )?); + let state = crate::state::AppState::new(pool.clone(), &config, search, "test-only".into()); + let edited: DateTime = "2026-10-01T12:00:00Z".parse()?; + let external_key = "notion://page"; + let old_body = "# Quarterly policy\n\nPolicy details\n"; + let complete_body = "# Quarterly policy\n\nPolicy details\nRevenue was 42 million.\n"; + + // Seed exactly the body the old top-level-only reader would have stored. + let created = ingest_item( + &state, + kb, + source, + external_key, + "Quarterly-policy.md", + "text/markdown", + old_body.as_bytes(), + Some(edited), + ) + .await?; + anyhow::ensure!(created == IngestAction::Created); + let original = documents::find_by_external_key(&pool, source, external_key) + .await? + .ok_or_else(|| anyhow::anyhow!("the original page was not persisted"))?; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/search")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "results": [{ + "id": "page", + "properties": {"Name": {"type": "title", "title": [{"plain_text": "Quarterly policy"}]}}, + "last_edited_time": "2026-10-01T12:00:00Z" + }], + "has_more": false + }))) + .expect(3) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/blocks/page/children")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "results": [{ + "id": "toggle", "type": "toggle", "has_children": true, + "toggle": {"rich_text": [{"plain_text": "Policy details"}]} + }], + "has_more": false + }))) + .expect(3) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/blocks/toggle/children")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "results": [{ + "id": "paragraph", "type": "paragraph", "has_children": false, + "paragraph": {"rich_text": [{"plain_text": "Revenue was 42 million."}]} + }], + "has_more": false + }))) + .expect(2) + .mount(&server) + .await; + let mut http = Paced::new(reqwest::Client::builder().no_proxy().build()?); + http.api_root = server.uri(); + + for expected in [IngestAction::Updated, IngestAction::Unchanged] { + let (pages, truncated, _) = fetch_pages(&mut http, None).await?; + anyhow::ensure!(!truncated && pages.len() == 1); + let page = &pages[0]; + anyhow::ensure!(page.external_key == external_key); + anyhow::ensure!(page.last_edited == Some(edited)); + anyhow::ensure!(page.text == complete_body); + let action = ingest_item( + &state, + kb, + source, + &page.external_key, + &page.filename, + "text/markdown", + page.text.as_bytes(), + page.last_edited, + ) + .await?; + anyhow::ensure!(action == expected, "expected {expected:?}, got {action:?}"); + let updated = documents::find_by_external_key(&pool, source, external_key) + .await? + .ok_or_else(|| anyhow::anyhow!("the resynced page was not persisted"))?; + anyhow::ensure!(updated.id == original.id); + anyhow::ensure!(updated.external_key.as_deref() == Some(external_key)); + anyhow::ensure!(updated.doc_time == original.doc_time); + anyhow::ensure!(updated.doc_time == Some(edited)); + anyhow::ensure!(updated.sha256 != original.sha256); + anyhow::ensure!(state.blob.get(&updated.sha256).await? == complete_body.as_bytes()); + anyhow::ensure!(state.blob.get(&original.sha256).await? == old_body.as_bytes()); + let versions: Vec<(i32, String)> = sqlx::query_as( + "SELECT version, sha256 FROM document_versions WHERE document_id=$1 ORDER BY version", + ) + .bind(updated.id) + .fetch_all(&pool) + .await?; + anyhow::ensure!( + versions == vec![(1, original.sha256.clone()), (2, updated.sha256.clone())], + "resync must retain the old body and create exactly one new version" + ); + if expected == IngestAction::Updated { + let jobs_before: i64 = sqlx::query_scalar( + "SELECT count(*) FROM jobs WHERE kind='process_document' AND payload->>'document_id'=$1", + ) + .bind(updated.id.to_string()) + .fetch_one(&pool) + .await?; + Mock::given(method("GET")) + .and(path("/blocks/toggle/children")) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "message": "temporarily unavailable" + }))) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + let (pages, _, unread) = fetch_pages(&mut http, None).await?; + anyhow::ensure!( + pages.is_empty(), + "a failed body read must not return a page to ingest" + ); + anyhow::ensure!( + unread + == vec!["notion page page: notion blocks returned 503 Service Unavailable: temporarily unavailable".to_string()] + ); + let retained = documents::find_by_external_key(&pool, source, external_key) + .await? + .ok_or_else(|| anyhow::anyhow!("the complete page was lost after a failed read"))?; + anyhow::ensure!(retained.sha256 == updated.sha256); + anyhow::ensure!(retained.updated_at == updated.updated_at); + anyhow::ensure!(state.blob.get(&retained.sha256).await? == complete_body.as_bytes()); + let versions_after: Vec<(i32, String)> = sqlx::query_as( + "SELECT version, sha256 FROM document_versions WHERE document_id=$1 ORDER BY version", + ) + .bind(updated.id) + .fetch_all(&pool) + .await?; + anyhow::ensure!(versions_after == versions); + let jobs_after: i64 = sqlx::query_scalar( + "SELECT count(*) FROM jobs WHERE kind='process_document' AND payload->>'document_id'=$1", + ) + .bind(updated.id.to_string()) + .fetch_one(&pool) + .await?; + anyhow::ensure!(jobs_after == jobs_before); + } + } + Ok::<_, anyhow::Error>(()) + } + .await; + sqlx::query("DELETE FROM knowledge_bases WHERE id=$1") + .bind(kb) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id=$1") + .bind(org) + .execute(&pool) + .await?; + result +} diff --git a/crates/utopia-server/src/notion_tests.rs b/crates/utopia-server/src/notion_tests.rs new file mode 100644 index 000000000..12b45f92f --- /dev/null +++ b/crates/utopia-server/src/notion_tests.rs @@ -0,0 +1,579 @@ +use super::*; +use serde_json::{json, Value}; +use std::sync::{Arc, Mutex}; +use wiremock::{matchers::method, matchers::path, Mock, MockServer, Request, ResponseTemplate}; + +fn paced(server: &MockServer) -> Paced { + // 本地 mock 不走开发机的系统代理。 + let mut http = Paced::new(reqwest::Client::builder().no_proxy().build().unwrap()); + http.api_root = server.uri(); + http +} + +fn block(id: &str, kind: &str, text: &str, has_children: bool) -> Value { + json!({ + "id": id, "type": kind, "has_children": has_children, + kind: {"rich_text": [{"plain_text": text}]} + }) +} + +async fn children( + server: &MockServer, + parent: &str, + cursor: Option<&str>, + blocks: Vec, + next_cursor: Option<&str>, +) { + let cursor = cursor.map(str::to_owned); + Mock::given(method("GET")) + .and(path(format!("/blocks/{parent}/children"))) + .and(move |request: &Request| { + let actual = request + .url + .query_pairs() + .find(|(key, _)| key == "start_cursor"); + actual.map(|(_, value)| value.into_owned()) == cursor + }) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "results": blocks, "has_more": next_cursor.is_some(), "next_cursor": next_cursor + }))) + .expect(1) + .mount(server) + .await; +} + +async fn request_paths(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() + .iter() + .map(|request| { + assert!(request + .url + .query_pairs() + .any(|(key, value)| key == "page_size" && value == "100")); + request.url.path().to_owned() + }) + .collect() +} + +#[tokio::test] +async fn toggle_body_is_read_before_the_next_sibling() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![ + block("toggle", "toggle", "Quarterly policy", true), + block("after", "paragraph", "After toggle", false), + ], + None, + ) + .await; + children( + &server, + "toggle", + None, + vec![block( + "paragraph", + "paragraph", + "Revenue was 42 million.", + false, + )], + None, + ) + .await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "Quarterly policy\nRevenue was 42 million.\nAfter toggle\n" + ); + assert_eq!( + request_paths(&server).await, + ["/blocks/page/children", "/blocks/toggle/children"] + ); +} + +#[tokio::test] +async fn columns_without_rich_text_keep_nested_body_in_reading_order() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![ + json!({"id":"columns", "type":"column_list", "has_children":true, "column_list":{}}), + block("after", "paragraph", "After columns", false), + ], + None, + ) + .await; + children( + &server, + "columns", + None, + vec![ + json!({"id":"left", "type":"column", "has_children":true, "column":{}}), + json!({"id":"right", "type":"column", "has_children":true, "column":{}}), + ], + None, + ) + .await; + children( + &server, + "left", + None, + vec![ + block("nested", "toggle", "Left toggle", true), + block("left-after", "paragraph", "Left after", false), + ], + None, + ) + .await; + children( + &server, + "nested", + None, + vec![block("body", "paragraph", "Deep body", false)], + None, + ) + .await; + children( + &server, + "right", + None, + vec![block("right-body", "paragraph", "Right body", false)], + None, + ) + .await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "Left toggle\nDeep body\nLeft after\nRight body\nAfter columns\n" + ); + assert_eq!( + request_paths(&server).await, + [ + "/blocks/page/children", + "/blocks/columns/children", + "/blocks/left/children", + "/blocks/nested/children", + "/blocks/right/children" + ] + ); +} + +#[tokio::test] +async fn pagination_cursors_belong_to_each_parent() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![ + block("a", "toggle", "A", true), + block("b", "toggle", "B", true), + ], + Some("root-next"), + ) + .await; + children( + &server, + "a", + None, + vec![block("a1", "paragraph", "A1", false)], + Some("a-next"), + ) + .await; + children( + &server, + "a", + Some("a-next"), + vec![block("a2", "paragraph", "A2", false)], + None, + ) + .await; + children( + &server, + "b", + None, + vec![block("b1", "paragraph", "B1", false)], + Some("b-next"), + ) + .await; + children( + &server, + "b", + Some("b-next"), + vec![block("b2", "paragraph", "B2", false)], + None, + ) + .await; + children( + &server, + "page", + Some("root-next"), + vec![block("end", "paragraph", "End", false)], + None, + ) + .await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "A\nA1\nA2\nB\nB1\nB2\nEnd\n" + ); + assert_eq!( + request_paths(&server).await, + [ + "/blocks/page/children", + "/blocks/a/children", + "/blocks/a/children", + "/blocks/b/children", + "/blocks/b/children", + "/blocks/page/children" + ] + ); +} + +#[tokio::test] +async fn block_budget_is_shared_and_stops_before_any_further_request() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![ + json!({"id":"container", "type":"column", "has_children":true, "column":{}}), + block("sibling", "toggle", "Must not appear", true), + ], + Some("root-next"), + ) + .await; + // 无文字容器也占一块；子层第五页的第 99 块耗尽整页预算。 + for page in 0..5 { + let cursor = format!("child-{page}"); + let next = format!("child-{}", page + 1); + let blocks = (0..100) + .map(|index| { + let id = format!("item-{}", page * 100 + index); + block(&id, "paragraph", &id, page == 4 && index == 98) + }) + .collect(); + children( + &server, + "container", + if page == 0 { None } else { Some(&cursor) }, + blocks, + Some(&next), + ) + .await; + } + let text = page_text(&mut paced(&server), "page").await.unwrap(); + assert_eq!(text.lines().count(), MAX_BLOCKS_PER_PAGE - 1); + assert!(text.ends_with("item-498\n")); + assert!(!text.contains("Must not appear")); + assert_eq!(request_paths(&server).await.len(), 6); +} + +#[tokio::test] +async fn exact_top_level_budget_does_not_fetch_another_page() { + let server = MockServer::start().await; + for page in 0..5 { + let cursor = format!("root-{page}"); + let next = format!("root-{}", page + 1); + let blocks = (0..100) + .map(|index| block(&format!("{page}-{index}"), "paragraph", "Body", false)) + .collect(); + children( + &server, + "page", + if page == 0 { None } else { Some(&cursor) }, + blocks, + Some(&next), + ) + .await; + } + assert_eq!( + page_text(&mut paced(&server), "page") + .await + .unwrap() + .lines() + .count(), + MAX_BLOCKS_PER_PAGE + ); + assert_eq!(request_paths(&server).await.len(), 5); +} + +#[tokio::test] +async fn independent_pages_and_databases_are_not_inlined() { + let server = MockServer::start().await; + children(&server, "page", None, vec![ + json!({"id":"child", "type":"child_page", "has_children":true, "child_page":{"title":"Separate page"}}), + json!({"id":"database", "type":"child_database", "has_children":true, "child_database":{"title":"Separate database"}}), + block("body", "paragraph", "Parent body", false), + ], None).await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "- Separate page\nParent body\n" + ); + assert_eq!(request_paths(&server).await, ["/blocks/page/children"]); +} + +#[tokio::test] +async fn flat_page_markdown_is_unchanged() { + let server = MockServer::start().await; + children(&server, "page", None, vec![ + block("heading", "heading_1", "Heading", false), + block("paragraph", "paragraph", "Body", false), + block("bullet", "bulleted_list_item", "Bullet", false), + block("number", "numbered_list_item", "Number", false), + json!({"type":"to_do", "to_do":{"checked":true, "rich_text":[{"plain_text":"Done"}]}}), + block("quote", "quote", "Quote", false), + json!({"type":"code", "code":{"language":"rust", "rich_text":[{"plain_text":"let x = 1;"}]}}), + json!({"type":"divider", "divider":{}}), + block("unknown", "new_type", "Unknown text", false), + ], None).await; + assert_eq!(page_text(&mut paced(&server), "page").await.unwrap(), + "## Heading\nBody\n- Bullet\n1. Number\n- [x] Done\n> Quote\n```rust\nlet x = 1;\n```\nUnknown text\n"); + assert_eq!(request_paths(&server).await.len(), 1); +} + +#[tokio::test] +async fn child_request_error_keeps_existing_diagnostics() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![block("child", "toggle", "Parent", true)], + None, + ) + .await; + Mock::given(path("/blocks/child/children")) + .respond_with( + ResponseTemplate::new(403).set_body_json(json!({"message":"child access denied"})), + ) + .expect(1) + .mount(&server) + .await; + let error = page_text(&mut paced(&server), "page").await.unwrap_err(); + assert_eq!( + error.to_string(), + "notion blocks returned 403 Forbidden: child access denied" + ); +} + +#[tokio::test] +async fn nested_requests_share_pacing_and_honor_retry_after() { + let server = MockServer::start().await; + let calls = Arc::new(Mutex::new(Vec::new())); + let recorded = calls.clone(); + Mock::given(method("GET")).respond_with(move |request: &Request| { + let mut calls = recorded.lock().unwrap(); + calls.push((request.url.path().to_owned(), Instant::now())); + match calls.len() { + 1 => ResponseTemplate::new(200).set_body_json(json!({"results":[block("child", "toggle", "Parent", true)], "has_more":false})), + 2 => ResponseTemplate::new(429).insert_header("Retry-After", "0.5").set_body_json(json!({"message":"slow down"})), + _ => ResponseTemplate::new(200).set_body_json(json!({"results":[block("body", "paragraph", "Child", false)], "has_more":false})), + } + }).expect(3).mount(&server).await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "Parent\nChild\n" + ); + let calls = calls.lock().unwrap(); + assert_eq!( + calls + .iter() + .map(|(path, _)| path.as_str()) + .collect::>(), + [ + "/blocks/page/children", + "/blocks/child/children", + "/blocks/child/children" + ] + ); + assert!(calls[1].1.duration_since(calls[0].1) >= MIN_INTERVAL); + assert!(calls[2].1.duration_since(calls[1].1) >= Duration::from_millis(500)); +} + +#[tokio::test] +async fn child_rate_limit_exhaustion_is_an_error() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![block("child", "toggle", "Parent", true)], + None, + ) + .await; + Mock::given(path("/blocks/child/children")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("Retry-After", "0") + .set_body_json(json!({"message":"still limited"})), + ) + .expect(u64::from(MAX_RATE_LIMIT_RETRIES + 1)) + .mount(&server) + .await; + let error = page_text(&mut paced(&server), "page").await.unwrap_err(); + assert_eq!( + error.to_string(), + "notion blocks returned 429 Too Many Requests: still limited" + ); +} + +#[tokio::test] +async fn a_page_whose_body_cannot_be_read_is_left_out_and_the_others_still_come_back() { + let server = MockServer::start().await; + let page = |id: &str, title: &str| { + json!({"id": id, "properties": {"Name": {"type": "title", "title": [{"plain_text": title}]}}, + "last_edited_time": "2026-10-01T12:00:00Z"}) + }; + Mock::given(method("POST")).and(path("/search")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "results": [page("page", "Quarterly policy"), page("other", "Travel policy")], "has_more": false + }))).expect(1).mount(&server).await; + children( + &server, + "page", + None, + vec![block("child", "toggle", "Partial body", true)], + None, + ) + .await; + Mock::given(path("/blocks/child/children")) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({"message":"unavailable"}))) + .expect(1) + .mount(&server) + .await; + children( + &server, + "other", + None, + vec![block( + "p", + "paragraph", + "Economy class under six hours", + false, + )], + None, + ) + .await; + + let (pages, truncated, unread) = fetch_pages(&mut paced(&server), None).await.unwrap(); + // 读坏的那一页不带着半截正文或只有标题回来，另一页不受它连累 + assert_eq!(pages.len(), 1); + assert_eq!(pages[0].external_key, "notion://other"); + assert!(pages[0].text.contains("Economy class under six hours")); + assert!(!truncated); + assert_eq!( + unread, + vec!["notion page page: notion blocks returned 503 Service Unavailable: unavailable"] + ); +} + +#[tokio::test] +async fn a_nested_table_keeps_its_header_cells_and_paginated_rows() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![block("toggle", "toggle", "Quarterly policy", true)], + None, + ) + .await; + children( + &server, + "toggle", + None, + vec![ + json!({"id":"table", "type":"table", "has_children":true, + "table":{"table_width":3, "has_column_header":true, "has_row_header":true}}), + block("after", "paragraph", "After table", false), + ], + None, + ) + .await; + children( + &server, + "table", + None, + vec![ + json!({"type":"table_row", "table_row":{"cells":[ + [{"plain_text":"Metric"}], [{"plain_text":"Value"}], [{"plain_text":"Notes"}] + ]}}), + json!({"type":"table_row", "table_row":{"cells":[ + [{"plain_text":"Revenue"}], [], [{"plain_text":"Pending"}] + ]}}), + ], + Some("table-next"), + ) + .await; + children( + &server, + "table", + Some("table-next"), + vec![json!({"type":"table_row", "table_row":{"cells":[ + [{"plain_text":"Pipe | slash \\"}], + [{"plain_text":"42"}, {"plain_text":" million"}], + [{"plain_text":"First\r\nSecond\nThird\rLast"}] + ]}})], + None, + ) + .await; + let text = page_text(&mut paced(&server), "page").await.unwrap(); + assert_eq!( + text, + "Quarterly policy\n\n| Metric | Value | Notes |\n| --- | --- | --- |\n| Revenue | | Pending |\n| Pipe \\| slash \\\\ | 42 million | First
Second
Third
Last |\n\nAfter table\n" + ); + assert_eq!( + request_paths(&server).await, + [ + "/blocks/page/children", + "/blocks/toggle/children", + "/blocks/table/children", + "/blocks/table/children" + ] + ); + let chunks = utopia_ingest::chunk_text(&text); + assert!(chunks + .iter() + .any(|chunk| chunk.text.contains("| --- | --- | --- |"))); +} + +#[tokio::test] +async fn a_table_without_column_headers_keeps_the_first_row_as_data() { + let server = MockServer::start().await; + children( + &server, + "page", + None, + vec![ + block("before", "paragraph", "Before table", false), + json!({"id":"table", "type":"table", "has_children":true, + "table":{"table_width":2, "has_column_header":false}}), + block("after", "paragraph", "After table", false), + ], + None, + ) + .await; + children( + &server, + "table", + None, + vec![ + json!({"type":"table_row", "table_row":{"cells":[ + [{"plain_text":"Revenue"}], [{"plain_text":"42 million"}] + ]}}), + json!({"type":"table_row", "table_row":{"cells":[[], []]}}), + ], + None, + ) + .await; + assert_eq!( + page_text(&mut paced(&server), "page").await.unwrap(), + "Before table\n\n| | |\n| --- | --- |\n| Revenue | 42 million |\n| | |\n\nAfter table\n" + ); +} diff --git a/crates/utopia-server/src/object_storage.rs b/crates/utopia-server/src/object_storage.rs index bb3253c42..7b6dffefe 100644 --- a/crates/utopia-server/src/object_storage.rs +++ b/crates/utopia-server/src/object_storage.rs @@ -381,15 +381,21 @@ mod tests { Ok(()) } - /// GCS 真连一次。**没有可用的模拟器，所以这条在 CI 上永远跳过。** + /// GCS against a real bucket. Run on 2026-10-05 in us-central1: flat + /// namespace, uniform access, a service account with only + /// `Storage Object Viewer` on the bucket. It passed. /// - /// 试过 `fsouza/fake-gcs-server`，不行：它只实现 JSON API - /// （`/storage/v1/b/...`），而 `object_store` 的 GCS 后端列对象走的是 - /// **XML API**（`/bucket?list-type=2`），那个路径上它回 404。 - /// 换句话说，最常见的那个 GCS 模拟器测不了我们实际用的那条协议路径。 + /// **CI still skips it: there is no usable emulator.** `fsouza/fake-gcs-server` + /// implements the JSON API (`/storage/v1/b/...`), while `object_store` lists + /// over the XML API (`/bucket?list-type=2`), where it answers 404. + /// + /// The bucket needs two things under `docs/`: + /// - `a.txt` containing `hello` + /// - one folder made in the Cloud Console. The console writes a zero-byte + /// placeholder (`docs/x/`) and `object_store` hands it back as `docs/x`, + /// the same shape that took S3 sync down. Without one the test cannot + /// reach that path, so it fails instead of passing hollow. /// - /// 所以 GCS 这条路径**只有构造测试覆盖**，跟 S3 与 Azure 不是一个成色。 - /// 有真桶的人跑一次就能补上： /// ```text /// UTOPIA_GCS_TEST_ENDPOINT=https://storage.googleapis.com \\ /// UTOPIA_GCS_TEST_BUCKET=your-bucket \\ @@ -410,11 +416,36 @@ mod tests { cfg["service_account_key"] = serde_json::Value::String(k); } let store = client("gcs", &cfg)?; + + // Look at the raw listing first. `fetch` drops placeholders by size, so + // its output alone cannot tell "filtered" from "never there". + let prefix = StorePath::from("docs"); + let mut raw = store.list(Some(&prefix)); + let mut placeholders = 0usize; + while let Some(meta) = raw.next().await { + if meta?.size == 0 { + placeholders += 1; + } + } + assert!( + placeholders > 0, + "no zero-byte placeholder under docs/: make a folder in the Cloud Console" + ); + let (objs, _) = fetch("gcs", store.as_ref(), &bucket, Some("docs")).await?; + let keys: Vec<&str> = objs.iter().map(|o| o.external_key.as_str()).collect(); + assert!( + objs.iter().all(|o| !o.bytes.is_empty()), + "a placeholder reached fetch's output: {keys:?}" + ); let a = objs.iter().find(|o| o.filename == "a.txt").expect("a.txt"); + assert_eq!(a.bytes, b"hello", "wrong content read back"); assert_eq!(a.external_key, format!("gs://{bucket}/docs/a.txt")); - assert!(a.last_modified.is_some()); + assert!( + a.last_modified.is_some(), + "LastModified is the only source of doc_time" + ); Ok(()) } } diff --git a/crates/utopia-server/src/phrase_alignment.rs b/crates/utopia-server/src/phrase_alignment.rs index e5e9d4c26..63160b2a9 100644 --- a/crates/utopia-server/src/phrase_alignment.rs +++ b/crates/utopia-server/src/phrase_alignment.rs @@ -289,7 +289,23 @@ async fn shortlist( return Ok(out); } let pool = &state.pool; - for batch in wide.chunks(SHORTLIST_EMBED_BATCH) { + // 先问这一类属性有没有向量，没有的签名不嵌：嵌了最近邻也是空，照样看全部。一个没跑过 + // `embed_ontology` 的库从前每轮把一万七千条签名送去嵌入、短名单一条没开（#1097）。 + // 一类问一句，至多两句 + let mut has_vectors = HashMap::new(); + for kind in ["relation", "attribute"] { + if wide.iter().any(|s| property_kind(s) == kind) { + let has = + utopia_store::ontology::has_relation_type_vectors(pool, kb_id, Some(kind)).await?; + has_vectors.insert(kind, has); + } + } + let usable: Vec<&PhraseSignature> = wide + .iter() + .copied() + .filter(|s| has_vectors.get(property_kind(s)) == Some(&true)) + .collect(); + for batch in usable.chunks(SHORTLIST_EMBED_BATCH) { let texts: Vec = batch .iter() .map(|s| match (s.examples.first(), s.quotes.first()) { @@ -314,17 +330,12 @@ async fn shortlist( }; for (s, vector) in batch.iter().zip(vectors) { let (fitting, _) = &full[&s.key()]; - let kind = if s.object_is_value { - "attribute" - } else { - "relation" - }; let near = utopia_store::ontology::nearest_relation_type_ids( pool, kb_id, &vector, (fitting.len() * 2) as i64, - Some(kind), + Some(property_kind(s)), ) .await?; if near.is_empty() { @@ -343,10 +354,25 @@ async fn shortlist( ); } } - tracing::info!(%kb_id, wide = wide.len(), shortlisted = out.len(), "候选短名单开好"); + tracing::info!( + %kb_id, + wide = wide.len(), + skipped = wide.len() - usable.len(), + shortlisted = out.len(), + "候选短名单开好" + ); Ok(out) } +/// 签名能绑的是哪一类属性：宾语是字面值的绑属性，否则绑关系 +fn property_kind(s: &PhraseSignature) -> &'static str { + if s.object_is_value { + "attribute" + } else { + "relation" + } +} + /// 类别词 → 留给提规则看的属性 id（按相关度）。词的文本加几个例名嵌入，取最近的 /// [`SHORTLIST`] 条；没配嵌入模型、属性没向量的不进表 async fn shortlist_kind_words( @@ -359,7 +385,10 @@ async fn shortlist_kind_words( let Some(client) = llm_util::embed_client(settings) else { return Ok(out); }; - if words.is_empty() { + // 一条带向量的属性都没有，嵌了也是空手而回（理由同 `shortlist`） + if words.is_empty() + || !utopia_store::ontology::has_relation_type_vectors(&state.pool, kb_id, None).await? + { return Ok(out); } for batch in words.chunks(SHORTLIST_EMBED_BATCH) { diff --git a/crates/utopia-server/src/phrase_alignment_tests.rs b/crates/utopia-server/src/phrase_alignment_tests.rs index cfa15275b..e01544b94 100644 --- a/crates/utopia-server/src/phrase_alignment_tests.rs +++ b/crates/utopia-server/src/phrase_alignment_tests.rs @@ -14,6 +14,8 @@ struct Model { hold: Arc, entered: Arc, release: Arc, + /// 嵌入端点被叫了几次：短名单该不该嵌签名（#1097）看的是这个数 + embeds: Arc, } async fn reply(State(m): State, Json(body): Json) -> impl IntoResponse { let n = { @@ -41,7 +43,8 @@ async fn reply(State(m): State, Json(body): Json) -> impl IntoResp /// 脚本化的嵌入端点：每段文字一个由字节算出的四维向量。只要求确定、条数对得上—— /// 这里测的是「向量有没有」，不是近不近 -async fn embed(Json(body): Json) -> impl IntoResponse { +async fn embed(State(m): State, Json(body): Json) -> impl IntoResponse { + m.embeds.fetch_add(1, std::sync::atomic::Ordering::SeqCst); let data: Vec = body["input"] .as_array() .cloned() @@ -130,6 +133,7 @@ impl Fx { hold: Arc::new(std::sync::atomic::AtomicBool::new(false)), entered: Arc::new(tokio::sync::Notify::new()), release: Arc::new(tokio::sync::Notify::new()), + embeds: Arc::new(std::sync::atomic::AtomicUsize::new(0)), }; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; let endpoint = format!("http://{}", listener.local_addr()?); @@ -251,6 +255,9 @@ impl Fx { fn requests(&self) -> Vec { self.model.requests.lock().unwrap().clone() } + fn embed_calls(&self) -> usize { + self.model.embeds.load(std::sync::atomic::Ordering::SeqCst) + } fn prompt_of(&self, n: usize) -> String { self.requests()[n]["messages"][1]["content"] .as_str() @@ -513,6 +520,44 @@ async fn overflow_is_recorded_for_a_person_and_recovers_when_candidates_shrink( run } +/// 属性还都没有向量时，短名单不嵌签名（#1097）：嵌了最近邻也是空，模型照样看全部结构 +/// 候选。从前一个没跑过 `embed_ontology` 的库每轮把候选多的签名全送去嵌入，一条也用不上。 +/// 属性有了向量，同一条签名照常嵌、照常开短名单 +#[tokio::test] +async fn signatures_are_not_embedded_while_no_property_has_a_vector() -> anyhow::Result<()> { + let Some(f) = Fx::new().await? else { + return Ok(()); + }; + let run = async { + f.enable_embeddings().await?; + // 12 条不声明域/值域的关系，加上 based_in 共 13 > SHORTLIST：这条签名要开短名单 + for i in 0..12 { + sqlx::query("INSERT INTO relation_types(id,kb_id,key,label,kind,temporal) VALUES($1,$2,$3,$3,'relation','state')") + .bind(Uuid::now_v7()).bind(f.kb).bind(format!("filler_{i}")).execute(&f.pool).await?; + } + f.script(bound()); + f.run().await?; + assert_eq!(f.embed_calls(), 0, "no property has a vector: nothing to embed for"); + assert_eq!(f.binding().await?.status, "bound"); + assert!( + f.prompt_of(0).contains("filler_11"), + "without a shortlist the model sees every fitting candidate" + ); + + sqlx::query("UPDATE relation_types SET embedding = '[1,0,0,0]' WHERE kb_id=$1") + .bind(f.kb) + .execute(&f.pool) + .await?; + f.script(bound()); + f.run().await?; + assert_eq!(f.embed_calls(), 1, "with property vectors the signature is embedded"); + anyhow::Ok(()) + } + .await; + f.cleanup().await?; + run +} + #[tokio::test] async fn an_endpoint_class_change_moves_the_signature_and_the_old_row_stops_looping( ) -> anyhow::Result<()> { diff --git a/crates/utopia-server/src/query_engine/clickhouse.rs b/crates/utopia-server/src/query_engine/clickhouse.rs new file mode 100644 index 000000000..88394abb3 --- /dev/null +++ b/crates/utopia-server/src/query_engine/clickhouse.rs @@ -0,0 +1,1293 @@ +//! ClickHouse over HTTP: `POST /` with the statement in the body. With no session, read-only +//! and the timeout are URL settings on each request, and **the account decides which it takes**: +//! +//! - **Drop a refused setting and retry.** Ask for `readonly=1` and `max_execution_time`; drop +//! the one the profile refuses (four measured account types: [`ClickHouseEngine::negotiated`]). +//! - **If the server cannot hold the time, the engine kills the query.** Each request has a +//! `query_id`; at the deadline, or when the caller drops the future (Stop in chat), we send +//! `KILL QUERY`. This needs `SELECT ON system.processes`; without it a query is refused, unless +//! readonly=1 holds and the account's own `max_execution_time` is 10 s or less. Behind a load +//! balancer with several replicas a KILL can reach another replica and match nothing; give such +//! accounts a server-side `max_execution_time`. +//! - **Values come back by column type** from `JSONCompactStrings`; [`cell`] says why. + +use super::conn::ClickHouseConn; +use super::{ + coerce, rows_to_json_lines, truncate_rows, wrap_limit, QueryEngine, QueryResult, SchemaColumn, + STATEMENT_TIMEOUT_SECS, +}; +use reqwest::Client; +use serde::Deserialize; +use std::sync::Arc; +use std::time::Duration; +use tokio::task::JoinHandle; + +/// How `JSONCompactStrings` writes NULL. When the server's `output_format_pretty_grid_charset` +/// is ASCII, it writes `NULL` +const NULL_TEXT: &str = "ᴺᵁᴸᴸ"; +const NULL_TEXT_ASCII: &str = "NULL"; + +pub struct ClickHouseEngine { + conn: Arc, + /// Kill the query if no reply comes by then. Two seconds more than the server's + /// `max_execution_time`: when both apply, the server's error is clearer than ours + deadline: Duration, +} + +#[derive(Deserialize)] +struct Reply { + meta: Vec, + data: Vec>, + /// An error raised after the result started. Servers 23.9 to 25.10 put it here by default + /// (`http_write_exception_in_output_format`), with status 200 and half the data. From 25.11 + /// the error text follows the output, and the body does not parse + exception: Option, +} + +#[derive(Deserialize)] +struct Meta { + name: String, + #[serde(rename = "type")] + type_name: String, +} + +/// An error from the server. `code` comes from the `X-ClickHouse-Exception-Code` header +#[derive(Debug, thiserror::Error)] +#[error("{message}")] +struct ServerError { + code: Option, + message: String, +} + +/// The two guards a request carries on its URL +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct Guards { + readonly: bool, + timeout: bool, +} + +impl Guards { + const ALL: Self = Self { + readonly: true, + timeout: true, + }; + const NONE: Self = Self { + readonly: false, + timeout: false, + }; +} + +#[derive(Debug, PartialEq, Eq)] +enum Refused { + Timeout, + ReadOnly, +} + +/// A ClickHouse string literal. ClickHouse also reads a backslash as an escape, so the shared +/// `sql_literal`, which only doubles quotes, would misread a name like `sales\north` +fn literal(s: &str) -> String { + format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) +} + +/// The last `n` characters: a long reply keeps its error at the end +fn tail(text: &str, n: usize) -> &str { + text.char_indices() + .rev() + .nth(n.saturating_sub(1)) + .map_or(text, |(i, _)| &text[i..]) +} + +/// The server's message from an error reply. An error after part of the result comes after +/// that data. reqwest asks for gzip, and when the server compresses (`enable_http_compression`, +/// on by default in 26.x) it holds back the headers up to a size: below it the reply is a 500 +/// with all rows before `Code: N.` (26.9.12.8: 15 MB); above it see [`trailer_message`]. The +/// message can quote the same `Code: N.` in a literal, so a reply that starts with it is the +/// plain error and stays whole +fn server_message(code: Option, text: &str) -> String { + let text = text.trim(); + if let Some(message) = trailer_message(text) { + return message.to_string(); + } + let Some(c) = code else { + return tail(text, 1000).to_string(); + }; + if text.starts_with(&format!("Code: {c}.")) { + return text.to_string(); + } + match text.rfind(&format!("Code: {c}. DB::")) { + Some(i) => text[i..].to_string(), + None => tail(text, 1000).to_string(), + } +} + +/// When the headers have already gone out (no compression, or a result too large to hold +/// back), 25.11+ sends status 200: the rows, then +/// `__exception__ __exception__`, one per line. The block must +/// end the body and open and close with the same tag +fn trailer_message(text: &str) -> Option<&str> { + let tag = trailer_tag(text)?; + let body = text + .strip_suffix("__exception__")? + .trim_end() + .rsplit_once('\n')? + .0; + let block = &body[body.rfind("__exception__")? + "__exception__".len()..]; + Some(block.trim_start().strip_prefix(tag)?.trim()) +} + +/// The tag on the last line of an `__exception__` block +fn trailer_tag(text: &str) -> Option<&str> { + let rest = text.strip_suffix("__exception__")?.trim_end(); + rest.rsplit_once('\n')?.1.split_whitespace().nth(1) +} + +fn query_id() -> String { + format!("utopia-{}", uuid::Uuid::now_v7().simple()) +} + +fn request( + client: &Client, + conn: &ClickHouseConn, + body: &str, + guards: Guards, + query_id: Option<&str>, +) -> anyhow::Result { + let mut url = reqwest::Url::parse(&conn.base)?; + { + let mut q = url.query_pairs_mut(); + if let Some(id) = query_id { + q.append_pair("query_id", id); + } + if let Some(db) = &conn.database { + q.append_pair("database", db); + } + if guards.readonly { + q.append_pair("readonly", "1"); + } + if guards.timeout { + q.append_pair("max_execution_time", &STATEMENT_TIMEOUT_SECS.to_string()); + } + } + let mut req = client.post(url).body(body.to_string()); + if let Some(user) = &conn.user { + req = req.basic_auth(user, conn.password.as_deref()); + } + Ok(req) +} + +/// Kill our own query. Gives back whether the KILL matched it: the reply has one TSV row per +/// matched query (status, query_id, user, query), and none when no query matched (for example, +/// the KILL reached another replica). The error holds the server's reason (usually no +/// `system.processes` grant) +async fn kill_query( + client: &Client, + conn: &ClickHouseConn, + query_id: &str, +) -> anyhow::Result { + let sql = format!( + "KILL QUERY WHERE query_id = {} ASYNC FORMAT TabSeparated", + literal(query_id) + ); + let resp = request(client, conn, &sql, Guards::NONE, None)? + .timeout(Duration::from_secs(5)) + .send() + .await?; + let status = resp.status(); + let text = resp.text().await?; + anyhow::ensure!(status.is_success(), "{}", text.trim()); + Ok(text + .lines() + .any(|row| row.split('\t').nth(1) == Some(query_id))) +} + +/// Kills the query if the future is dropped (Stop in chat, MCP disconnect) or at the deadline +/// ([`KillOnDrop::fire`]). The KILL runs in its own task, so a later drop cannot stop it +struct KillOnDrop { + target: Option<(Client, Arc, String)>, +} + +impl KillOnDrop { + fn disarm(&mut self) { + self.target = None; + } + + fn fire(mut self) -> Option>> { + self.target.take().and_then(spawn_kill) + } +} + +impl Drop for KillOnDrop { + fn drop(&mut self) { + if let Some(target) = self.target.take() { + spawn_kill(target); + } + } +} + +fn spawn_kill( + (client, conn, query_id): (Client, Arc, String), +) -> Option>> { + let rt = tokio::runtime::Handle::try_current().ok()?; + Some(rt.spawn(async move { + let killed = kill_query(&client, &conn, &query_id).await; + match &killed { + Err(e) => tracing::warn!(%query_id, error = %e, "could not kill the ClickHouse query"), + Ok(false) => tracing::warn!(%query_id, "KILL QUERY matched no ClickHouse query"), + Ok(true) => {} + } + killed + })) +} + +impl ClickHouseEngine { + pub fn new(conn: ClickHouseConn) -> Self { + Self { + conn: Arc::new(conn), + deadline: Duration::from_secs(u64::from(STATEMENT_TIMEOUT_SECS) + 2), + } + } + + async fn post(&self, client: &Client, body: &str, guards: Guards) -> anyhow::Result { + let id = query_id(); + let req = request(client, &self.conn, body, guards, Some(&id))?; + let mut guard = KillOnDrop { + target: Some((client.clone(), self.conn.clone(), id)), + }; + let sent = tokio::time::timeout(self.deadline, async { + let mut resp = req.send().await?; + let status = resp.status(); + let header = |name| { + resp.headers() + .get(name) + .and_then(|v| v.to_str().ok()) + .map(|v| v.trim().to_string()) + }; + let code = header("X-ClickHouse-Exception-Code").and_then(|v| v.parse::().ok()); + let tag = header("X-ClickHouse-Exception-Tag"); + let mut body = Vec::new(); + loop { + match resp.chunk().await { + Ok(Some(chunk)) => body.extend_from_slice(&chunk), + Ok(None) => break, + // After an error block 25.11+ cuts the reply on purpose, so that it cannot pass + // for a whole result. A block with the header's tag is the server's answer; + // any other cut is a transport error + Err(e) => { + let text = String::from_utf8_lossy(&body); + match (&tag, trailer_tag(text.trim())) { + (Some(want), Some(got)) if want == got => break, + _ => return Err(e), + } + } + } + } + let text = String::from_utf8(body) + .unwrap_or_else(|e| String::from_utf8_lossy(e.as_bytes()).into_owned()); + Ok::<_, reqwest::Error>((status, code, text)) + }) + .await; + let Ok(sent) = sent else { + let secs = self.deadline.as_secs(); + let killed = match guard.fire() { + Some(task) => task.await.unwrap_or_else(|e| Err(e.into())), + None => Err(anyhow::anyhow!("no runtime to send KILL QUERY from")), + }; + return Err(match killed { + Ok(true) => anyhow::anyhow!( + "ClickHouse query did not finish within {secs}s and was cancelled" + ), + Ok(false) => anyhow::anyhow!( + "ClickHouse query did not finish within {secs}s and could not be cancelled \ + (KILL QUERY matched no query; behind a load balancer it may have reached \ + another replica); it may still be running" + ), + Err(e) => anyhow::anyhow!( + "ClickHouse query did not finish within {secs}s and could not be cancelled \ + ({e}); it may still be running" + ), + }); + }; + let (status, code, text) = match sent { + // A success, or an error with ClickHouse's exception header, means ClickHouse answered + // and the statement has ended. A reply without that header (a gateway's 502 or 504) + // proves nothing, so the guard stays armed + Ok(reply) => { + if reply.0.is_success() || reply.1.is_some() { + guard.disarm(); + } + reply + } + // Only a connect error proves the statement never reached the server. After any other + // error (say, the connection drops while we read the reply) it may still run + Err(e) => { + if e.is_connect() { + guard.disarm(); + } + return Err(e.into()); + } + }; + if !status.is_success() { + return Err(ServerError { + code, + message: server_message(code, &text), + } + .into()); + } + Ok(text) + } + + /// Sends with read-only and the timeout, drops a refused setting and sends again, and returns + /// the guards it finally sent. What each account type accepts (measured on 26.9.12.8): + /// + /// - normal account: both, and the server enforces both. Use `readonly=1`, not 2: with + /// `readonly=2`, `SELECT … SETTINGS max_execution_time = 0` turns the timeout off; + /// - profile with `readonly=1`: `max_execution_time` gets READONLY (164); `readonly=1` is + /// the current value, so the server accepts it; + /// - profile with `readonly=2`: the reverse (`readonly=1` refused); statements can change it; + /// - `max_execution_time` locked by a constraint (`CONST`): SETTING_CONSTRAINT_VIOLATION (452). + /// + /// Drop `readonly` only if the account itself is read-only: a statement with `SETTINGS + /// readonly = 0` gets the same "Cannot modify 'readonly'", and dropping it would open the gate + async fn negotiated(&self, client: &Client, body: &str) -> anyhow::Result<(String, Guards)> { + let mut guards = Guards::ALL; + loop { + if guards != Guards::ALL { + self.can_stop_in_time(client, guards).await?; + } + let err = match self.post(client, body, guards).await { + Ok(text) => return Ok((text, guards)), + Err(e) => e, + }; + let refused = err + .downcast_ref::() + .and_then(|e| refused_setting(e.code, &e.message)); + match refused { + Some(Refused::Timeout) if guards.timeout => guards.timeout = false, + Some(Refused::ReadOnly) if guards.readonly => { + // If the probe fails or the account itself is 0, return the original refusal + let own = self.account_setting(client, "readonly").await; + if !matches!(own.as_deref(), Ok("1" | "2")) { + return Err(err); + } + guards.readonly = false; + } + _ => return Err(err), + } + } + } + + /// One setting of the account itself (from its profile), asked without settings. FORMAT is + /// fixed: a profile can set `default_format` to JSONEachRow, which returns a JSON object + async fn account_setting(&self, client: &Client, name: &str) -> anyhow::Result { + let sql = format!("SELECT getSetting({}) FORMAT TabSeparated", literal(name)); + Ok(self + .post(client, &sql, Guards::NONE) + .await? + .trim() + .to_string()) + } + + /// Runs before a statement goes out without one of the guards. The server holds the time + /// only with both guards, so otherwise only KILL is left: try it now on an unused id, before + /// the statement runs. If KILL fails, an account running under readonly=1 (a profile or our + /// own setting) may rely on its own max_execution_time, including a CONST one, but + /// only if it is not longer than ours: a 3600 s limit does not hold our 10 s + async fn can_stop_in_time(&self, client: &Client, guards: Guards) -> anyhow::Result<()> { + // A fresh id never matches, so only the permission counts here + let Err(reason) = kill_query(client, &self.conn, &query_id()).await else { + return Ok(()); + }; + let limit = if guards.readonly { + let own = self.account_setting(client, "max_execution_time").await; + own.ok().and_then(|t| t.parse::().ok()) + } else { + None + }; + if limit.is_some_and(|t| t > 0.0 && t <= f64::from(STATEMENT_TIMEOUT_SECS)) { + return Ok(()); + } + anyhow::bail!( + "This ClickHouse account cannot cancel its own queries ({reason}), so nothing would \ + stop a query at the {STATEMENT_TIMEOUT_SECS}s timeout. Grant it SELECT ON \ + system.processes, or set max_execution_time to at most {STATEMENT_TIMEOUT_SECS} in \ + its read-only settings profile." + ) + } + + async fn run(&self, client: &Client, sql: &str) -> anyhow::Result<(Reply, Guards)> { + let (text, guards) = self + .negotiated(client, &format!("{sql}\nFORMAT JSONCompactStrings")) + .await?; + let reply: Reply = + serde_json::from_str(&text).map_err(|_| match trailer_message(text.trim()) { + Some(message) => anyhow::anyhow!("{message}"), + None => anyhow::anyhow!( + "ClickHouse returned an unreadable result: {}", + tail(&text, 300).trim() + ), + })?; + if let Some(e) = &reply.exception { + anyhow::bail!("{}", e.trim()); + } + Ok((reply, guards)) + } +} + +/// Which of our settings the error refuses: `Cannot modify 'X' setting in readonly mode` (164) +/// for a read-only account, `Setting X should not be changed` (452) for a constraint. A write +/// that read-only blocks is also 164, but its message names no setting +fn refused_setting(code: Option, message: &str) -> Option { + if !matches!(code, Some(164) | Some(452)) { + return None; + } + let names = |name: &str| { + message.contains(&format!("'{name}' setting")) + || message.contains(&format!("Setting {name} ")) + }; + if names("max_execution_time") { + Some(Refused::Timeout) + } else if names("readonly") { + Some(Refused::ReadOnly) + } else { + None + } +} + +/// Strips `LowCardinality(...)` / `Nullable(...)` (they can nest); the bool means "can be NULL" +fn unwrap_type(type_name: &str) -> (&str, bool) { + let mut ty = type_name.trim(); + let mut nullable = false; + loop { + if let Some(inner) = ty + .strip_prefix("Nullable(") + .and_then(|t| t.strip_suffix(')')) + { + ty = inner; + nullable = true; + } else if let Some(inner) = ty + .strip_prefix("LowCardinality(") + .and_then(|t| t.strip_suffix(')')) + { + ty = inner; + } else { + return (ty, nullable); + } + } +} + +/// One cell: a `JSONCompactStrings` string becomes a JSON value. +/// +/// Not `JSONCompact`: quotes on 64-bit and wider integers depend on the server's +/// `output_format_json_quote_64bit_integers` (no quotes by default since 25.8, quotes before), +/// and a read-only account cannot change it. serde_json reads an unquoted 128 / 256-bit integer +/// as f64 and rounds it. Not `coerce` either: its number list has Snowflake's `FIXED`, so the +/// zip code `'02134'` in a `FixedString(5)` becomes 2134. Only these types convert: +/// +/// - integers: a number if it fits i64 / u64, else (128 / 256 bits) an exact string. An exact +/// string is better for the model than a number with wrong last digits; +/// - `Decimal(P, S)`: a number if P ≤ 15 (f64 holds it exactly), else a string, same reason; +/// - floats become numbers (`nan` / `inf` stay strings: not JSON numbers), `Bool` a bool; +/// - NULL only in `Nullable(...)` columns, and the ASCII `NULL` only in non-string columns. A +/// `Nullable(String)` cell that holds that text is ambiguous: the cost of exact integers +fn cell(type_name: &str, raw: &serde_json::Value) -> serde_json::Value { + use serde_json::Value; + let (base, nullable) = unwrap_type(type_name); + let Value::String(s) = raw else { + return raw.clone(); + }; + let text_type = base == "String" + || ["FixedString(", "Enum8(", "Enum16("] + .iter() + .any(|t| base.starts_with(t)); + if nullable && (s == NULL_TEXT || (s == NULL_TEXT_ASCII && !text_type)) { + return Value::Null; + } + let integer = base + .strip_prefix("UInt") + .or_else(|| base.strip_prefix("Int")) + .is_some_and(|bits| !bits.is_empty() && bits.chars().all(|c| c.is_ascii_digit())); + if integer { + if let Ok(n) = s.parse::() { + return n.into(); + } + if let Ok(n) = s.parse::() { + return n.into(); + } + return raw.clone(); + } + let short_decimal = base + .strip_prefix("Decimal(") + .and_then(|rest| rest.split(',').next()) + .and_then(|p| p.trim().parse::().ok()) + .is_some_and(|p| p <= 15); + if short_decimal || matches!(base, "Float32" | "Float64" | "BFloat16") { + return coerce("DOUBLE", raw); + } + if base == "Bool" { + return coerce("BOOL", raw); + } + raw.clone() +} + +#[async_trait::async_trait] +impl QueryEngine for ClickHouseEngine { + async fn test(&self) -> anyhow::Result<()> { + self.run(&super::http()?, "SELECT 1").await.map(|_| ()) + } + + async fn fetch_schema(&self) -> anyhow::Result> { + // With a database in the URL, read only that one: the default database is often empty, and + // without a scope the whole server goes into the schema document. Keep `type` as is, with + // `Nullable(...)` / `LowCardinality(...)`: the first says NULLs occur, the second usually + // marks a dimension column. Both help, and the SQL is the same as for the inner type + let scope = match &self.conn.database { + Some(db) => format!("database = {}", literal(db)), + None => "database NOT IN ('system', 'INFORMATION_SCHEMA', 'information_schema')".into(), + }; + let client = super::http()?; + let (reply, _) = self + .run( + &client, + &format!( + "SELECT database, table, name, type, comment FROM system.columns \ + WHERE {scope} ORDER BY database, table, position" + ), + ) + .await?; + Ok(reply + .data + .into_iter() + .map(super::trino::schema_row) + .collect()) + } + + async fn execute(&self, sql: &str) -> anyhow::Result { + let client = super::http()?; + let (reply, _) = self.run(&client, &wrap_limit(sql)).await?; + let columns: Vec = reply.meta.iter().map(|m| m.name.clone()).collect(); + let rows: Vec> = reply + .data + .iter() + .map(|row| { + row.iter() + .zip(&reply.meta) + .map(|(v, m)| cell(&m.type_name, v)) + .collect() + }) + .collect(); + let (rows, truncated) = truncate_rows(rows); + Ok(QueryResult { + rows: rows_to_json_lines(&columns, &rows), + truncated, + }) + } +} + +#[cfg(test)] +mod live_tests { + use super::super::conn::ClickHouseConn; + use super::super::QueryEngine; + use super::ClickHouseEngine; + use std::time::{Duration, Instant}; + + /// The real-server test; it skips without `UTOPIA_TEST_CLICKHOUSE_URL` (measured on 26.9.12.8): + /// `docker run -d -p 18124:8123 -e CLICKHOUSE_PASSWORD=pw -e + /// CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT=1 clickhouse/clickhouse-server:26.9.12.8`, run the SQL + /// below as default, then once per account, e.g. `clickhouse://ro:pw@127.0.0.1:18124/default`: + /// ```sql + /// CREATE TABLE orders (id UInt64, region LowCardinality(String), note Nullable(String), + /// amount Decimal(12, 2) COMMENT 'CNY', big UInt64) ENGINE = MergeTree ORDER BY id; + /// INSERT INTO orders VALUES (1, 'east', NULL, 1234.56, 18446744073709551615); + /// CREATE USER ro IDENTIFIED BY 'pw' SETTINGS readonly = 1; + /// CREATE USER ro2 IDENTIFIED BY 'pw' SETTINGS readonly = 2; + /// CREATE USER pinned IDENTIFIED BY 'pw' SETTINGS max_execution_time = 30 CONST; + /// CREATE USER nokill IDENTIFIED BY 'pw' SETTINGS readonly = 1; + /// GRANT SELECT, INSERT, CREATE TABLE ON default.* TO ro, ro2, pinned, nokill; + /// GRANT SELECT ON system.processes TO ro, ro2, pinned; -- the engine needs it; nokill has none + /// ``` + fn live_url(var: &str) -> Option { + std::env::var(var).ok().filter(|u| !u.trim().is_empty()) + } + + #[tokio::test] + async fn a_live_server_answers_read_only_and_within_the_timeout() { + let Some(url) = live_url("UTOPIA_TEST_CLICKHOUSE_URL") else { + return; + }; + let engine = ClickHouseEngine::new(ClickHouseConn::parse(&url).expect("parse url")); + engine.test().await.expect("test()"); + + let schema = engine.fetch_schema().await.expect("schema"); + let col = |name: &str| { + schema + .iter() + .find(|c| c.table == "orders" && c.column == name) + .unwrap_or_else(|| panic!("orders.{name}")) + }; + assert_eq!(col("region").data_type, "LowCardinality(String)"); + assert_eq!(col("note").data_type, "Nullable(String)"); + assert_eq!(col("amount").comment.as_deref(), Some("CNY")); + assert_eq!( + col("note").comment, + None, + "an empty comment is filtered out" + ); + assert!(!col("id").is_primary_key, "a sort key is not a unique key"); + + // The UInt64 maximum would become 18446744073709552000 through f64 + let r = engine + .execute("SELECT id, region, note, amount, big FROM orders ORDER BY id") + .await + .expect("execute"); + assert_eq!( + r.rows, + vec![ + r#"{"id":1,"region":"east","note":null,"amount":1234.56,"big":18446744073709551615}"# + ] + ); + + // The server refuses the write: the normal account through our readonly=1, the others + // through their own settings. This bypasses the gate to prove layer 3 + let client = super::super::http().unwrap(); + let err = engine + .negotiated( + &client, + "CREATE TABLE orders_probe (x UInt8) ENGINE = Memory", + ) + .await + .expect_err("a write must be refused") + .to_string(); + assert!(err.contains("readonly mode"), "{err}"); + + // Timeout: the server stops the query for the normal and readonly=2 accounts, our KILL + // for the rest. Both must return near the deadline and leave no query on the server + let started = Instant::now(); + let err = engine + .execute("SELECT sum(sipHash64(number)) FROM numbers(1000000000000)") + .await + .expect_err("the slow query must be stopped") + .to_string(); + let took = started.elapsed(); + assert!( + err.contains("TIMEOUT_EXCEEDED") || err.contains("did not finish"), + "{err}" + ); + assert!( + took >= Duration::from_secs(9) && took < engine.deadline + Duration::from_secs(3), + "took {took:?}: {err}" + ); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let (left, _) = engine + .run( + &client, + "SELECT count() FROM system.processes WHERE \ + query LIKE '%numbers(1000000000000)%' AND query NOT LIKE '%system.processes%'", + ) + .await + .expect("processes"); + if left.data[0][0] == serde_json::json!("0") { + break; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("the slow query is still on the server: {err}")); + } + + /// An account that cannot kill its queries and has no short profile limit is refused, on Test + /// connection and on every query. Point `UTOPIA_TEST_CLICKHOUSE_NO_KILL_URL` at `nokill` above + #[tokio::test] + async fn an_account_that_cannot_kill_its_queries_fails_the_test() { + let Some(url) = live_url("UTOPIA_TEST_CLICKHOUSE_NO_KILL_URL") else { + return; + }; + let engine = ClickHouseEngine::new(ClickHouseConn::parse(&url).expect("parse url")); + let err = engine + .test() + .await + .expect_err("test() must refuse") + .to_string(); + assert!(err.contains("SELECT ON system.processes"), "{err}"); + let err = engine + .execute("SELECT 1") + .await + .expect_err("a query must be refused without test() too") + .to_string(); + assert!(err.contains("SELECT ON system.processes"), "{err}"); + } +} + +#[cfg(test)] +mod tests { + use super::super::conn::ClickHouseConn; + use super::super::{QueryEngine, QueryResult}; + use super::{cell, literal, refused_setting, unwrap_type, ClickHouseEngine, Refused}; + use serde_json::json; + use std::time::Duration; + use wiremock::matchers::{ + body_string_contains, method, query_param, query_param_contains, query_param_is_missing, + }; + use wiremock::{Mock, MockBuilder, MockServer, ResponseTemplate}; + + const REFUSED_TIMEOUT: &str = + "Cannot modify 'max_execution_time' setting in readonly mode. (READONLY)"; + const REFUSED_READONLY: &str = "Cannot modify 'readonly' setting in readonly mode. (READONLY)"; + const SLOW: &str = "SELECT sum(number) AS s FROM numbers(1000000000000)"; + + fn engine(server: &MockServer) -> ClickHouseEngine { + let host = server.uri().trim_start_matches("http://").to_string(); + ClickHouseEngine::new( + ClickHouseConn::parse(&format!("clickhouse://u:p@{host}/sales?ssl=false")).unwrap(), + ) + } + + fn mock_post() -> MockBuilder { + Mock::given(method("POST")) + } + + /// Each mock must get exactly one request + async fn once(server: &MockServer, mock: MockBuilder, response: ResponseTemplate) { + mock.respond_with(response).expect(1).mount(server).await; + } + + fn reply(meta: &[(&str, &str)], data: serde_json::Value) -> ResponseTemplate { + let meta: Vec<_> = meta + .iter() + .map(|(n, t)| json!({ "name": n, "type": t })) + .collect(); + ResponseTemplate::new(200).set_body_json(json!({ "meta": meta, "data": data, "rows": 1 })) + } + + fn refusal(code: u32, message: &str) -> ResponseTemplate { + ResponseTemplate::new(500) + .insert_header("X-ClickHouse-Exception-Code", code.to_string().as_str()) + .set_body_string(format!("Code: {code}. DB::Exception: {message}")) + } + + fn text(body: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_string(format!("{body}\n")) + } + + #[test] + fn values_come_back_by_column_type() { + let i128_min = "-170141183460469231731687303715884105728"; + let decimal38 = "12345678901234567890.0123456789"; + for (ty, raw, want) in [ + ( + "UInt64", + "18446744073709551615", + json!(18446744073709551615u64), + ), + ("Int64", "-5", json!(-5)), + // Wider than 64 bits stays an exact string, not an approximation + ("Int128", i128_min, json!(i128_min)), + ("Nullable(Int32)", "ᴺᵁᴸᴸ", json!(null)), + ("Nullable(Int32)", "NULL", json!(null)), + ("LowCardinality(Nullable(String))", "ᴺᵁᴸᴸ", json!(null)), + // Text NULL in a string column is ambiguous, so it stays; non-Nullable means no NULL + ("Nullable(String)", "NULL", json!("NULL")), + ("String", "ᴺᵁᴸᴸ", json!("ᴺᵁᴸᴸ")), + // An Enum label is text too: a label `NULL` stays, the real marker is null + ( + "Nullable(Enum8('NULL' = 1, 'z' = 2))", + "NULL", + json!("NULL"), + ), + ("Nullable(Enum8('NULL' = 1, 'z' = 2))", "ᴺᵁᴸᴸ", json!(null)), + // Leading zeros are part of the value + ("FixedString(5)", "02134", json!("02134")), + ("Decimal(12, 2)", "1234.56", json!(1234.56)), + ("Decimal(38, 10)", decimal38, json!(decimal38)), + ("Nullable(Bool)", "true", json!(true)), + ("Float64", "1.5", json!(1.5)), + ("Float64", "nan", json!("nan")), + ("Date", "2023-06-01", json!("2023-06-01")), + ("IntervalDay", "3", json!("3")), + ("Array(UInt8)", "[1,2]", json!("[1,2]")), + ] { + assert_eq!(cell(ty, &json!(raw)), want, "{ty} {raw}"); + } + } + + /// ClickHouse reads a backslash as an escape: `'sales\north'` would name `salesorth` + #[test] + fn a_literal_escapes_backslashes_and_quotes() { + assert_eq!(literal(r"sales\north's"), r"'sales\\north\'s'"); + } + + #[tokio::test] + async fn the_schema_filter_keeps_a_backslash_in_the_database_name() { + let server = MockServer::start().await; + let host = server.uri().trim_start_matches("http://").to_string(); + let filter = mock_post().and(body_string_contains(r"database = 'sales\\north'")); + once(&server, filter, reply(&[("database", "String")], json!([]))).await; + let conn = ClickHouseConn::parse(&format!("clickhouse://u@{host}/sales%5Cnorth")).unwrap(); + ClickHouseEngine::new(conn).fetch_schema().await.unwrap(); + } + + #[test] + fn wrappers_come_off_in_any_order() { + for (ty, want) in [ + ("LowCardinality(Nullable(String))", ("String", true)), + ("Nullable(Decimal(12, 2))", ("Decimal(12, 2)", true)), + ("LowCardinality(String)", ("String", false)), + ("Array(Nullable(UInt8))", ("Array(Nullable(UInt8))", false)), + ] { + assert_eq!(unwrap_type(ty), want, "{ty}"); + } + } + + #[test] + fn only_a_refused_setting_counts_as_one() { + let constraint = + "Setting max_execution_time should not be changed. (SETTING_CONSTRAINT_VIOLATION)"; + for (code, message, want) in [ + (164, REFUSED_TIMEOUT, Some(Refused::Timeout)), + (164, REFUSED_READONLY, Some(Refused::ReadOnly)), + (452, constraint, Some(Refused::Timeout)), + // A write that read-only blocks: also 164, but no setting was refused + ( + 164, + "rw: Cannot execute query in readonly mode. (READONLY)", + None, + ), + (62, "Syntax error: 'readonly' setting", None), + ] { + assert_eq!(refused_setting(Some(code), message), want, "{message}"); + } + } + + #[tokio::test] + async fn sends_read_only_and_a_timeout_and_keeps_column_order() { + let server = MockServer::start().await; + let request = mock_post() + .and(query_param("readonly", "1")) + .and(query_param("max_execution_time", "10")) + .and(query_param("database", "sales")) + .and(body_string_contains("LIMIT 201")) + .and(body_string_contains("FORMAT JSONCompactStrings")); + let meta = [ + ("region", "LowCardinality(String)"), + ("total", "UInt64"), + ("note", "Nullable(String)"), + ]; + let data = json!([["east", "18446744073709551615", "ᴺᵁᴸᴸ"]]); + once(&server, request, reply(&meta, data)).await; + let out = engine(&server) + .execute("SELECT region, total, note FROM orders") + .await + .unwrap(); + assert_eq!( + out.rows, + vec![r#"{"region":"east","total":18446744073709551615,"note":null}"#] + ); + assert!(!out.truncated); + } + + /// An error after the result started: valid JSON, status 200, half the data + #[tokio::test] + async fn an_exception_inside_the_result_is_an_error() { + let server = MockServer::start().await; + let half = json!({ + "meta": [{ "name": "number", "type": "UInt64" }], + "data": [["0"]], + "exception": "Code: 396. DB::Exception: Limit for result exceeded. (TOO_MANY_ROWS_OR_BYTES)" + }); + once( + &server, + mock_post(), + ResponseTemplate::new(200).set_body_json(half), + ) + .await; + let err = engine(&server).execute(SLOW).await.unwrap_err().to_string(); + assert!(err.contains("TOO_MANY_ROWS_OR_BYTES"), "{err}"); + } + + /// The same error with a gzip reply: status 500, and the rows come before the message. + /// The message quotes `Code: 395.` once more, as ClickHouse does with a literal + #[tokio::test] + async fn an_error_after_rows_keeps_only_the_message() { + let server = MockServer::start().await; + let rows = format!( + r#"{{"meta":[{{"name":"s","type":"String"}}],"data":[["{}"]"#, + "x".repeat(100_000) + ); + let message = "Code: 395. DB::Exception: boom: while executing \ + 'throwIf(..., 'Code: 395. boom'_String)'. (FUNCTION_THROW_IF_VALUE_IS_NON_ZERO)"; + let reply = ResponseTemplate::new(500) + .insert_header("X-ClickHouse-Exception-Code", "395") + .set_body_string(format!("{rows}{message}\n")); + once(&server, mock_post(), reply).await; + let err = engine(&server).execute(SLOW).await.unwrap_err().to_string(); + assert_eq!(err, message); + } + + /// When the headers have already gone out, the same error comes as status 200 with an + /// `__exception__` block + #[tokio::test] + async fn an_error_block_after_rows_keeps_only_the_message() { + let server = MockServer::start().await; + let message = "Code: 395. DB::Exception: boom. (FUNCTION_THROW_IF_VALUE_IS_NON_ZERO)"; + let body = format!( + "{{\n\t\"meta\": [{{\"name\": \"s\", \"type\": \"String\"}}],\n\t\"data\":\n\t[\n\t\t[\"0\"]\r\n\ + __exception__\r\nabcdef\r\n{message}\n{} abcdef\r\n__exception__\r\n", + message.len() + 1 + ); + once( + &server, + mock_post(), + ResponseTemplate::new(200).set_body_string(body), + ) + .await; + let err = engine(&server).execute(SLOW).await.unwrap_err().to_string(); + assert_eq!(err, message); + } + + /// A long ordinary error keeps its start, which names the cause + #[tokio::test] + async fn a_long_error_stays_whole() { + let server = MockServer::start().await; + let message = format!( + "Code: 47. DB::Exception: Unknown expression identifier `missing_col`. In scope {}. \ + (UNKNOWN_IDENTIFIER)", + "SELECT a AS b, ".repeat(200) + ); + once( + &server, + mock_post(), + ResponseTemplate::new(500) + .insert_header("X-ClickHouse-Exception-Code", "47") + .set_body_string(message.clone()), + ) + .await; + let err = engine(&server).execute(SLOW).await.unwrap_err().to_string(); + assert_eq!(err, message); + } + + /// A readonly=1 profile: max_execution_time is refused and dropped, readonly=1 stays. + /// The server no longer holds the time, so `test()` also tries a KILL + #[tokio::test] + async fn a_refused_timeout_is_dropped_and_read_only_stays() { + let server = MockServer::start().await; + let first = mock_post().and(query_param("max_execution_time", "10")); + once(&server, first, refusal(164, REFUSED_TIMEOUT)).await; + let retry = mock_post() + .and(query_param_is_missing("max_execution_time")) + .and(query_param("readonly", "1")); + once(&server, retry, reply(&[("1", "UInt8")], json!([["1"]]))).await; + let kill = mock_post().and(body_string_contains("KILL QUERY")); + once(&server, kill, ResponseTemplate::new(200)).await; + engine(&server).test().await.unwrap(); + } + + /// A readonly=1 account that cannot kill its queries, queried with `execute()` alone (no + /// `test()`). `profile_limit` is its profile's `max_execution_time` (`None`: the probe fails). + /// The retry without the timeout must go out `retries` times + async fn query_without_kill( + profile_limit: Option<&str>, + retries: u64, + ) -> anyhow::Result { + let server = MockServer::start().await; + let first = mock_post().and(query_param("max_execution_time", "10")); + once(&server, first, refusal(164, REFUSED_TIMEOUT)).await; + let no_grant = "Not enough privileges. To execute this query, it's necessary to have \ + the grant SELECT ON system.processes. (ACCESS_DENIED)"; + let kill = mock_post().and(body_string_contains("KILL QUERY")); + once(&server, kill, refusal(497, no_grant)).await; + let probe = mock_post().and(body_string_contains("getSetting('max_execution_time')")); + let limit = profile_limit.map_or_else(|| refusal(497, "(ACCESS_DENIED)"), text); + once(&server, probe, limit).await; + mock_post() + .and(query_param_is_missing("max_execution_time")) + .and(body_string_contains("SELECT 7 AS x")) + .respond_with(reply(&[("x", "UInt8")], json!([["7"]]))) + .expect(retries) + .mount(&server) + .await; + engine(&server).execute("SELECT 7 AS x").await + } + + /// No KILL, and the profile limit cannot hold 10 s (none, too long, unknown): the query is + /// refused with the grant message, and the retry without the timeout never goes out + #[tokio::test] + async fn a_query_that_nothing_can_stop_in_time_is_refused() { + for limit in [Some("0"), Some("3600"), None] { + let err = query_without_kill(limit, 0).await.unwrap_err().to_string(); + assert!( + err.contains("Grant it SELECT ON system.processes"), + "{limit:?}: {err}" + ); + assert!(err.contains("ACCESS_DENIED"), "{limit:?}: {err}"); + } + } + + /// No KILL, but the profile limit is 10 s or less: the server holds it, so the query runs + #[tokio::test] + async fn a_short_profile_limit_stands_in_for_kill() { + for limit in ["5", "10"] { + let out = query_without_kill(Some(limit), 1).await.unwrap(); + assert_eq!(out.rows, vec![r#"{"x":7}"#], "{limit}"); + } + } + + /// readonly=1 is refused (a readonly=2 profile, or a statement's own `SETTINGS readonly = 0`), + /// and the account's own readonly is `own`. Only a read-only account drops readonly and retries + async fn readonly_refused(own: &str) -> anyhow::Result { + let server = MockServer::start().await; + let probe = mock_post().and(body_string_contains( + "getSetting('readonly') FORMAT TabSeparated", + )); + once(&server, probe, text(own)).await; + let first = mock_post().and(query_param("readonly", "1")); + once(&server, first, refusal(164, REFUSED_READONLY)).await; + // A read-only account resends the same statement without readonly, after a KILL check + // (statements can lift its timeout); a writable one must not + let retries = if own == "0" { 0 } else { 1 }; + mock_post() + .and(body_string_contains("KILL QUERY")) + .respond_with(ResponseTemplate::new(200)) + .expect(retries) + .mount(&server) + .await; + mock_post() + .and(query_param_is_missing("readonly")) + .and(query_param("max_execution_time", "10")) + .and(body_string_contains("SELECT 7 AS x")) + .and(body_string_contains("LIMIT 201")) + .respond_with(reply(&[("x", "UInt8")], json!([["7"]]))) + .expect(retries) + .mount(&server) + .await; + engine(&server).execute("SELECT 7 AS x").await + } + + #[tokio::test] + async fn read_only_is_dropped_only_for_an_account_that_is_read_only_itself() { + let out = readonly_refused("2") + .await + .expect("a readonly=2 account retries without readonly=1"); + assert_eq!(out.rows, vec![r#"{"x":7}"#]); + let err = readonly_refused("0") + .await + .expect_err("a writable account keeps readonly=1 and the refusal") + .to_string(); + assert!(err.contains("Cannot modify 'readonly'"), "{err}"); + } + + /// A slow query that answers after 5 s, plus a KILL endpoint (no timeout, exactly one hit) + /// `matched`: the KILL reply has the row ClickHouse gives for a matched query, or is empty + async fn slow_server(matched: bool) -> MockServer { + let server = MockServer::start().await; + mock_post() + .and(body_string_contains("KILL QUERY WHERE query_id = 'utopia-")) + .and(query_param_is_missing("max_execution_time")) + .respond_with(move |req: &wiremock::Request| { + let body = String::from_utf8_lossy(&req.body); + let id = body.split('\'').nth(1).unwrap_or_default(); + let row = format!("waiting\t{id}\tu\tSELECT ...\n"); + ResponseTemplate::new(200).set_body_string(if matched { + row + } else { + String::new() + }) + }) + .expect(1) + .mount(&server) + .await; + let slow = reply(&[("s", "UInt64")], json!([["1"]])).set_delay(Duration::from_secs(5)); + mock_post() + .and(query_param_contains("query_id", "utopia-")) + .and(body_string_contains("numbers")) + .respond_with(slow) + .mount(&server) + .await; + server + } + + /// The KILL names the slow query's `query_id`. The guard sends it from another task, so wait + async fn assert_killed_the_slow_query(server: &MockServer) { + let requests = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let requests = server.received_requests().await.unwrap_or_default(); + if requests.len() >= 2 { + break requests; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("no KILL was sent"); + let body = |r: &wiremock::Request| String::from_utf8_lossy(&r.body).into_owned(); + let slow = requests + .iter() + .find(|r| body(r).contains("numbers")) + .and_then(|r| r.url.query_pairs().find(|(k, _)| k == "query_id")) + .map(|(_, v)| v.into_owned()) + .expect("the slow query carries a query_id"); + let kill = requests + .iter() + .map(body) + .find(|b| b.contains("KILL QUERY")) + .expect("a KILL was sent"); + assert!(kill.contains(&format!("'{slow}'")), "KILL names {slow}"); + } + + /// No reply by the deadline: a timeout error, and the query is killed + #[tokio::test] + async fn a_query_past_the_deadline_is_killed() { + let server = slow_server(true).await; + let mut ch = engine(&server); + ch.deadline = Duration::from_millis(300); + let err = ch.execute(SLOW).await.unwrap_err().to_string(); + assert!( + err.contains("did not finish") && err.contains("was cancelled"), + "{err}" + ); + assert_killed_the_slow_query(&server).await; + } + + /// The KILL matched nothing (say, it reached another replica): not cancelled + #[tokio::test] + async fn a_kill_that_matches_nothing_is_not_reported_as_cancelled() { + let server = slow_server(false).await; + let mut ch = engine(&server); + ch.deadline = Duration::from_millis(300); + let err = ch.execute(SLOW).await.unwrap_err().to_string(); + assert!( + err.contains("could not be cancelled") && err.contains("may still be running"), + "{err}" + ); + } + + /// The caller stops waiting (Stop in chat): the future is dropped, and the guard still kills it + #[tokio::test] + async fn a_dropped_query_is_killed_too() { + let server = slow_server(true).await; + let ch = engine(&server); + let stopped = tokio::time::timeout(Duration::from_millis(300), ch.execute(SLOW)).await; + assert!(stopped.is_err(), "the query should still be running"); + assert_killed_the_slow_query(&server).await; + } + + /// Reads one HTTP request (head and body) from a raw socket + async fn read_request(sock: &mut tokio::net::TcpStream) -> String { + use tokio::io::AsyncReadExt; + let (mut buf, mut chunk) = (Vec::new(), [0u8; 4096]); + loop { + let n = sock.read(&mut chunk).await.unwrap_or(0); + buf.extend_from_slice(&chunk[..n]); + let text = String::from_utf8_lossy(&buf).into_owned(); + let body_len = text + .lines() + .find_map(|l| { + l.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(str::to_owned) + }) + .and_then(|v| v.trim().parse::().ok()) + .unwrap_or(0); + match text.find("\r\n\r\n") { + Some(end) if buf.len() >= end + 4 + body_len => return text, + _ if n == 0 => return text, + _ => {} + } + } + } + + /// The server takes the statement, then the connection drops while we read the reply. The + /// query may still run, so the guard kills it + #[tokio::test] + async fn a_reply_cut_short_kills_the_query() { + use tokio::io::AsyncWriteExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (seen, mut requests) = tokio::sync::mpsc::unbounded_channel::(); + tokio::spawn(async move { + while let Ok((mut sock, _)) = listener.accept().await { + let request = read_request(&mut sock).await; + let reply: &[u8] = if request.contains("KILL QUERY") { + b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n" + } else { + b"HTTP/1.1 200 OK\r\ncontent-length: 100\r\n\r\n{\"meta\"" + }; + let _ = sock.write_all(reply).await; + let _ = seen.send(request); + } + }); + let conn = ClickHouseConn::parse(&format!("clickhouse://u@{addr}/sales")).unwrap(); + let err = ClickHouseEngine::new(conn).execute(SLOW).await.unwrap_err(); + assert!(!err.to_string().contains("did not finish"), "{err}"); + let wait = Duration::from_secs(5); + let statement = tokio::time::timeout(wait, requests.recv()) + .await + .expect("the statement") + .unwrap(); + let id = statement + .split(['?', '&', ' ']) + .find_map(|p| p.strip_prefix("query_id=")) + .expect("the statement carries a query_id") + .to_string(); + let kill = tokio::time::timeout(wait, requests.recv()) + .await + .expect("a KILL was sent") + .unwrap(); + assert!( + kill.contains("KILL QUERY") && kill.contains(&format!("'{id}'")), + "{kill}" + ); + } + + /// A refused connection never reached the server, so there is nothing to kill + #[tokio::test] + async fn a_refused_connection_sends_no_kill() { + let port = std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port(); + let conn = + ClickHouseConn::parse(&format!("clickhouse://u@127.0.0.1:{port}/sales")).unwrap(); + let err = ClickHouseEngine::new(conn).execute(SLOW).await.unwrap_err(); + assert!( + format!("{err:?}").to_lowercase().contains("connect"), + "{err:?}" + ); + // This test runs on one thread, so a KILL the guard spawned has not run yet. Open the + // port now, before the next await, and listen for it + let listener = std::net::TcpListener::bind(("127.0.0.1", port)).unwrap(); + listener.set_nonblocking(true).unwrap(); + let listener = tokio::net::TcpListener::from_std(listener).unwrap(); + let knock = tokio::time::timeout(Duration::from_millis(500), listener.accept()).await; + assert!(knock.is_err(), "a KILL was sent"); + } + + /// A gateway in front of ClickHouse gives up and answers a full 502 without ClickHouse's + /// header: the query may still run, so the guard kills it. A ClickHouse error reply (with + /// the header) means the statement has ended, so nothing is killed + #[tokio::test] + async fn a_gateway_error_kills_the_query_and_a_clickhouse_error_does_not() { + for (answer, kills) in [ + (ResponseTemplate::new(502).set_body_string("Bad Gateway"), 1), + (refusal(62, "Syntax error. (SYNTAX_ERROR)"), 0), + ] { + let server = MockServer::start().await; + mock_post() + .and(body_string_contains("KILL QUERY")) + .respond_with(ResponseTemplate::new(200)) + .expect(kills) + .mount(&server) + .await; + let slow = mock_post().and(body_string_contains("numbers")); + once(&server, slow, answer).await; + engine(&server).execute(SLOW).await.unwrap_err(); + if kills == 1 { + assert_killed_the_slow_query(&server).await; + } else { + // Give a wrongly spawned KILL time to arrive before the mock checks its count + tokio::time::sleep(Duration::from_millis(300)).await; + } + } + } +} diff --git a/crates/utopia-server/src/query_engine/conn.rs b/crates/utopia-server/src/query_engine/conn.rs index 1d3cd5192..0fd8f59c3 100644 --- a/crates/utopia-server/src/query_engine/conn.rs +++ b/crates/utopia-server/src/query_engine/conn.rs @@ -1,9 +1,10 @@ -//! 连接串解析。一个输入框、四种 scheme；这里把 URL 拆成各引擎要的字段。 +//! 连接串解析。一个输入框，scheme 决定引擎；这里把 HTTP 族的 URL 拆成各引擎要的字段。 //! //! 写法沿用 `postgres://user:pass@host/db` 的形状：凭据在 userinfo 里，HTTP 族的 //! 令牌放 password 位（`databricks://:TOKEN@…`），路径是「目录 / 库 / schema」， -//! 引擎特有的开关走 query。`ssl=false` 让 HTTP 族走明文——给本地代理与测试用， -//! 线上的三家都只认 https。 +//! 引擎特有的开关走 query。`ssl=false` 让 HTTP 族走明文——给本地代理与测试用。 +//! Databricks 与 Snowflake 线上只认 https；Trino 与 ClickHouse 默认明文，按端口与 +//! `ssl` 升级（各自的规则见 `TrinoConn` / `ClickHouseConn`）。 use percent_encoding::percent_decode_str; use url::Url; @@ -93,6 +94,37 @@ impl TrinoConn { } } +/// `clickhouse://user[:password]@host[:port]/[database][?ssl=true|false]` +/// +/// 走 HTTP 接口：明文 8123，https 8443。`ssl=true`（clickhouse-go 的 `secure=true` 也认） +/// 或端口 443 / 8443 时走 https，`ssl=false` 强制明文。跟 Trino 不同，带密码**不**升级 +/// https：ClickHouse 在明文口上照收密码，自建实例多半就是这样。不写用户就是服务器的 +/// `default` 用户；库名是之后每条查询的默认库，也是取 schema 的范围。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClickHouseConn { + pub base: String, + pub user: Option, + pub password: Option, + pub database: Option, +} + +impl ClickHouseConn { + pub fn parse(conn: &str) -> anyhow::Result { + let u = Url::parse(conn.trim())?; + let https = !ssl_off(&u) + && (query(&u, "ssl").as_deref() == Some("true") + || query(&u, "secure").as_deref() == Some("true") + || matches!(u.port(), Some(443) | Some(8443))); + let base = base_of(&u, https, if https { 8443 } else { 8123 })?; + Ok(Self { + base, + user: Some(decode(u.username())).filter(|s| !s.is_empty()), + password: u.password().map(decode).filter(|s| !s.is_empty()), + database: segments(&u).first().cloned(), + }) + } +} + /// `databricks://:TOKEN@workspace-host/sql/1.0/warehouses/WAREHOUSE_ID[?catalog=main&schema=default]` /// /// 路径就是 JDBC 里的 httpPath，从控制台复制过来不用改；`?warehouse=ID` 也认。 @@ -179,7 +211,34 @@ impl SnowflakeConn { #[cfg(test)] mod tests { - use super::{DatabricksConn, SnowflakeConn, TrinoConn}; + use super::{ClickHouseConn, DatabricksConn, SnowflakeConn, TrinoConn}; + + #[test] + fn clickhouse_is_plain_http_on_8123_unless_asked_for_tls() { + let c = ClickHouseConn::parse("clickhouse://analyst:p%40ss@ch.internal/sales").unwrap(); + assert_eq!(c.base, "http://ch.internal:8123"); + assert_eq!(c.user.as_deref(), Some("analyst")); + // 带密码不升级：这一点跟 Trino 相反，见 ClickHouseConn 上的注释 + assert_eq!(c.password.as_deref(), Some("p@ss")); + assert_eq!(c.database.as_deref(), Some("sales")); + + let c = ClickHouseConn::parse("clickhouse://u:p@abc.clickhouse.cloud:8443/db").unwrap(); + assert_eq!(c.base, "https://abc.clickhouse.cloud:8443"); + let c = ClickHouseConn::parse("clickhouse://u@ch.internal?ssl=true").unwrap(); + assert_eq!(c.base, "https://ch.internal:8443"); + assert_eq!(c.database, None); + let c = ClickHouseConn::parse("clickhouse://u@ch.internal:9443?secure=true").unwrap(); + assert_eq!(c.base, "https://ch.internal:9443"); + let c = ClickHouseConn::parse("clickhouse://u@127.0.0.1:8443?ssl=false").unwrap(); + assert_eq!(c.base, "http://127.0.0.1:8443"); + + // 没有用户就交给服务器的 default 用户 + let c = ClickHouseConn::parse("clickhouse://ch.internal:8123").unwrap(); + assert_eq!(c.user, None); + assert_eq!(c.password, None); + + assert!(ClickHouseConn::parse("clickhouse:///sales").is_err()); + } #[test] fn trino_defaults_to_plain_http_and_upgrades_when_it_must() { diff --git a/crates/utopia-server/src/query_engine/mod.rs b/crates/utopia-server/src/query_engine/mod.rs index a8097cf31..532352ee2 100644 --- a/crates/utopia-server/src/query_engine/mod.rs +++ b/crates/utopia-server/src/query_engine/mod.rs @@ -4,18 +4,21 @@ //! （`mysql.rs`，一条协议顺带覆盖 TiDB / OceanBase / Doris / StarRocks / MariaDB） //! → HTTP 族—— //! `trino.rs` 一个顶起 Iceberg / Delta / Hive 整个湖仓生态，`databricks.rs`、 -//! `snowflake.rs` 各走自家的 SQL REST API。挂载模型与注册表引擎无关，加引擎只放宽 -//! 一条 CHECK。连接串是唯一的输入：引擎由 scheme 决定（[`engine_from_conn`]）， -//! 剩下的部分各引擎自己拆（`conn.rs`），凭据只在服务端流转。 +//! `snowflake.rs` 各走自家的 SQL REST API，`clickhouse.rs` 走 ClickHouse 的 HTTP 接口。 +//! 挂载模型与注册表引擎无关，加引擎只放宽一条 CHECK。连接串是唯一的输入：引擎由 +//! scheme 决定（[`engine_from_conn`]），剩下的部分各引擎自己拆（`conn.rs`）， +//! 凭据只在服务端流转。 //! //! 安全闸（纵深防御，不信任模型）： //! 1. sqlparser 解析：仅放行单条 SELECT/WITH（含 CTE），拒绝 DML/DDL/多语句/SELECT INTO。 //! 按引擎选方言；sqlparser 没有 Trino 方言，Generic 是它的超集 //! 2. 强制外包一层 LIMIT（cap+1 探测截断） //! 3. 会话级只读 + 语句超时（引擎各自的机制，parser 万一漏网也写不进去）。 -//! HTTP 族没有会话，只有语句超时——只读靠第 1 层，这是它们比线协议少的那一层 +//! HTTP 族没有会话，只有语句超时——只读靠第 1 层，这是它们比线协议少的那一层。 +//! ClickHouse 例外：每个请求自己带 `readonly`，见 `clickhouse.rs` //! 4. 结果统一为 JSON Lines：PG 让库自己转；HTTP 族拿到列名与值后在这里拼，列序保留 +mod clickhouse; mod conn; mod databricks; mod mysql; @@ -25,7 +28,8 @@ mod trino; use sqlparser::ast::Statement; use sqlparser::dialect::{ - DatabricksDialect, GenericDialect, MySqlDialect, PostgreSqlDialect, SnowflakeDialect, + ClickHouseDialect, DatabricksDialect, GenericDialect, MySqlDialect, PostgreSqlDialect, + SnowflakeDialect, }; use sqlparser::parser::Parser; use std::time::Duration; @@ -38,7 +42,14 @@ pub(crate) const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(20); pub(crate) const HTTP_POLL_BUDGET: Duration = Duration::from_secs(30); /// 注册表里 `engine` 列的取值。迁移里的 CHECK 与这张表要一致 -pub const ENGINES: &[&str] = &["postgres", "mysql", "trino", "databricks", "snowflake"]; +pub const ENGINES: &[&str] = &[ + "postgres", + "mysql", + "trino", + "databricks", + "snowflake", + "clickhouse", +]; #[derive(Debug)] pub struct QueryResult { @@ -56,8 +67,8 @@ pub struct SchemaColumn { pub comment: Option, /// 该列本身是这张表的主键（单列主键；组合主键里这一位仍为 false）。 /// Postgres（#671）与 MySQL / MariaDB（#678）从 catalog 读出来；Trino / Snowflake / - /// Databricks 恒为 false，意思是「不知道」，不是「不是键」（各自能不能读见 - /// `trino::schema_row` 上的注释）。 + /// Databricks / ClickHouse 恒为 false，意思是「不知道」，不是「不是键」 + /// （各自能不能读见 `trino::schema_row` 上的注释）。 /// 探索提示词靠它把 ID 与量分开，宽表上一个八十列的 schema 没有这个 /// 几乎认不出哪一列是键（#502） pub is_primary_key: bool, @@ -77,7 +88,7 @@ pub trait QueryEngine: Send + Sync { async fn execute(&self, sql: &str) -> anyhow::Result; } -/// scheme → 引擎名。界面只有一个连接串输入框，这里是它唯一的分派点。 +/// scheme → 引擎名。界面上每档引擎的表单最后都拼成一条连接串（`dsSpecs`），这里是它唯一的分派点。 pub fn engine_from_conn(conn: &str) -> Option<&'static str> { let scheme = conn.trim().split("://").next()?.to_ascii_lowercase(); match scheme.as_str() { @@ -88,6 +99,7 @@ pub fn engine_from_conn(conn: &str) -> Option<&'static str> { "trino" | "presto" => Some("trino"), "databricks" => Some("databricks"), "snowflake" => Some("snowflake"), + "clickhouse" => Some("clickhouse"), _ => None, } } @@ -106,6 +118,9 @@ pub fn engine_for(engine: &str, conn: &str) -> anyhow::Result Ok(Box::new(snowflake::SnowflakeEngine::new( conn::SnowflakeConn::parse(conn)?, ))), + "clickhouse" => Ok(Box::new(clickhouse::ClickHouseEngine::new( + conn::ClickHouseConn::parse(conn)?, + ))), other => anyhow::bail!("Unsupported engine: {other}"), } } @@ -119,6 +134,7 @@ pub fn guard_sql_for(engine: &str, sql: &str) -> anyhow::Result { let parsed = match engine { "databricks" => Parser::parse_sql(&DatabricksDialect {}, cleaned), "snowflake" => Parser::parse_sql(&SnowflakeDialect {}, cleaned), + "clickhouse" => Parser::parse_sql(&ClickHouseDialect {}, cleaned), "trino" => Parser::parse_sql(&GenericDialect {}, cleaned), "mysql" => Parser::parse_sql(&MySqlDialect {}, cleaned), _ => Parser::parse_sql(&PostgreSqlDialect {}, cleaned), @@ -149,7 +165,7 @@ fn statement_kind(s: &Statement) -> &'static str { } } -/// 第 2 层：外包一层 LIMIT。三个 HTTP 引擎都认这个写法；PG 有自己的 row_to_json 版本 +/// 第 2 层：外包一层 LIMIT。HTTP 族的引擎都认这个写法；PG 有自己的 row_to_json 版本 pub(crate) fn wrap_limit(sql: &str) -> String { format!("SELECT * FROM ( {sql} ) AS _q LIMIT {}", ROW_CAP + 1) } @@ -185,7 +201,9 @@ pub(crate) fn rows_to_json_lines( .collect() } -/// Databricks 的 JSON_ARRAY 与 Snowflake 的 data 把每个值都给成字符串（或 null）。 +/// Databricks 的 JSON_ARRAY 与 Snowflake 的 data 把每个值都给成字符串（或 null）， +/// ClickHouse 的 JSONCompactStrings 也是（它只把浮点与布尔交给这里， +/// 见 `clickhouse::cell`）。 /// 按列类型把数与布尔还原，其余留字符串——模型对 `"42"` 和 `42` 的算术不一样 pub(crate) fn coerce(type_name: &str, raw: &serde_json::Value) -> serde_json::Value { let serde_json::Value::String(s) = raw else { @@ -319,7 +337,14 @@ mod tests { #[test] fn every_dialect_keeps_the_same_gate() { - for engine in ["postgres", "mysql", "trino", "databricks", "snowflake"] { + for engine in [ + "postgres", + "mysql", + "trino", + "databricks", + "snowflake", + "clickhouse", + ] { assert!( guard_sql_for(engine, "SELECT a FROM t WHERE b > 1").is_ok(), "{engine}" @@ -336,6 +361,14 @@ mod tests { assert!(guard_sql_for("trino", "SELECT count(*) FROM hive.default.orders").is_ok()); // MySQL 的反引号与 PG 方言不兼容，走自己的方言才过得去 assert!(guard_sql_for("mysql", "SELECT `region` FROM `sales`.`orders`").is_ok()); + // ClickHouse 的函数名与 FINAL 修饰都要过得去。语句里的 SETTINGS 闸这一层不管： + // readonly=1 生效时服务器拒掉它；readonly=2 的账号不拒，那时超时靠 KILL + // （见 clickhouse.rs） + assert!(guard_sql_for( + "clickhouse", + "SELECT toStartOfMonth(placed_on) AS m, count() FROM sales.orders FINAL GROUP BY m" + ) + .is_ok()); } #[test] @@ -355,6 +388,10 @@ mod tests { assert_eq!(engine_from_conn("mysql://u:p@h:3306/db"), Some("mysql")); // 同一套协议的另一个写法，引擎里会被改写成 mysql:// 再交给驱动 assert_eq!(engine_from_conn("mariadb://u@h/db"), Some("mysql")); + assert_eq!( + engine_from_conn("clickhouse://u:p@h:8123/sales"), + Some("clickhouse") + ); assert_eq!(engine_from_conn("garbage"), None); } diff --git a/crates/utopia-server/src/query_engine/trino.rs b/crates/utopia-server/src/query_engine/trino.rs index 8736b00e8..771a86973 100644 --- a/crates/utopia-server/src/query_engine/trino.rs +++ b/crates/utopia-server/src/query_engine/trino.rs @@ -179,14 +179,17 @@ pub(crate) fn schema_row(row: Vec) -> SchemaColumn { .and_then(|v| v.as_str()) .filter(|s| !s.is_empty()) .map(str::to_string), - // 键在这里一律「不知道」（#502）。三个引擎各不相同，别照 SQL 标准去读： + // 键在这里一律「不知道」（#502）。四个引擎各不相同，别照 SQL 标准去读： // - Trino / Presto：每个 catalog 的 `information_schema` 都由内建连接器提供，只有 // columns、tables、views、schemata 与权限、角色几张表，**没有** // `table_constraints` / `key_column_usage`——查它们是 TABLE_NOT_FOUND（#682）； // - Snowflake：没有 `key_column_usage`，列级的键要 `SHOW PRIMARY KEYS` / // `SHOW IMPORTED KEYS`； // - Databricks：Unity Catalog 在 `information_schema.table_constraints` / - // `key_column_usage` / `referential_constraints` 里登记信息性的主外键，读得到 + // `key_column_usage` / `referential_constraints` 里登记信息性的主外键，读得到； + // - ClickHouse：`system.columns` 有 `is_in_primary_key`，但主键是排序键的前缀， + // 不保证唯一——MergeTree 照收主键相同的两行，标成键会让探索提示词拿一个 + // 会重复的列当 ID is_primary_key: false, references_table: None, } diff --git a/crates/utopia-server/src/rdf.rs b/crates/utopia-server/src/rdf.rs index a0e20a055..01ffbf561 100644 --- a/crates/utopia-server/src/rdf.rs +++ b/crates/utopia-server/src/rdf.rs @@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex}; use chrono::{DateTime, Utc}; use oxrdf::vocab::{rdf, rdfs, xsd}; -use oxrdf::{Literal, NamedNode, NamedNodeRef, Term, TripleRef}; +use oxrdf::{Literal, NamedNode, Term, TripleRef}; use utopia_store::export::{ ExportAxiomViolation, ExportClass, ExportDerived, ExportDocument, ExportEntity, ExportFact, ExportFactConflict, ExportRelation, @@ -809,14 +809,128 @@ fn is_relative(v: &serde_json::Value) -> bool { /// 相对的值写成普通字符串：`"45 days after the Trigger Date"^^xsd:date` 是个不合法的字面量 fn literal_value(v: &serde_json::Value, datatype: Option<&str>) -> Option { let (text, prose) = literal_text(v)?; - let ty: NamedNodeRef<'_> = match datatype { - _ if prose || is_relative(v) => xsd::STRING, - Some("number") => xsd::DECIMAL, - Some("date") => xsd::DATE, - Some("bool") => xsd::BOOLEAN, - _ => xsd::STRING, + if prose || is_relative(v) { + return Some(Literal::new_simple_literal(text)); + } + Some(match datatype { + Some("number") => Literal::new_typed_literal(decimal_text(text), xsd::DECIMAL), + Some("date") => { + date_literal(&text).unwrap_or_else(|| Literal::new_typed_literal(text, xsd::DATE)) + } + Some("bool") => Literal::new_typed_literal(text, xsd::BOOLEAN), + _ => Literal::new_simple_literal(text), + }) +} + +/// JSON 数可以带指数，xsd:decimal 不可以。用已有的精确十进制类型展开文本， +/// 不再过一次 f64，也不舍入；没有指数的写法（包括负零）原样保留。 +fn decimal_text(text: String) -> String { + let Some((coefficient, exponent)) = text.split_once(['e', 'E']) else { + return text; + }; + // 归一化后的有限 f64 指数只在这个范围内。其他字符串不一定经过归一化， + // 不让异常的大指数触发无界分配，也不顺手清洗 BigDecimal 能容忍的下划线等写法。 + let supported_exponent = exponent + .parse::() + .is_ok_and(|exponent| (-324..=308).contains(&exponent)); + let decimal_coefficient = coefficient + .chars() + .all(|c| c.is_ascii_digit() || matches!(c, '+' | '-' | '.')); + if supported_exponent && decimal_coefficient { + if let Ok(number) = text.parse::() { + return number.to_plain_string(); + } + } + text +} + +/// 属性日期保留的是原文，不是世界时间区间。parse_time 会截掉小数秒、把无时区 +/// 钟点降到天，所以只在没有钟点时借它判断年/月/日；时钟和时区从原文规范化。 +/// 无法识别的文本由调用方沿用原输出；输入校验仍可能接受非日期的钟点文本， +/// 本次不收紧它的规则。年份按字段格式化，不能按四个字符切片；chrono 在扩展 +/// 正年份前加的 `+` 也需去掉，因为 XSD 年份不允许正号。 +fn date_literal(value: &str) -> Option { + let Some((date, clock)) = value.split_once(['T', ' ']) else { + let (date, precision) = utopia_extract::parse_time(value)?; + let (format, datatype) = match precision { + "year" => ("%Y", xsd::G_YEAR), + "month" => ("%Y-%m", xsd::G_YEAR_MONTH), + _ => ("%Y-%m-%d", xsd::DATE), + }; + let lexical = date.format(format).to_string(); + return Some(Literal::new_typed_literal( + lexical.trim_start_matches('+'), + datatype, + )); + }; + let date = chrono::NaiveDate::parse_from_str(date, "%Y-%m-%d").ok()?; + let calendar_date = date.format("%Y-%m-%d").to_string(); + let calendar_date = calendar_date.trim_start_matches('+'); + let (clock, zone) = if let Some(clock) = clock.strip_suffix(['Z', 'z']) { + (clock, "Z".to_string()) + } else if let Some(index) = clock.find(['+', '-']) { + let (clock, zone) = clock.split_at(index); + (clock, xsd_timezone(zone)?) + } else { + // XSD 允许没有时区的 dateTime；不能替原文猜一个 UTC。 + (clock, String::new()) + }; + let parts: Vec<&str> = clock.split(':').collect(); + let (hour, minute, second) = match parts.as_slice() { + [hour] => (*hour, "00", "00"), + [hour, minute] => (*hour, *minute, "00"), + [hour, minute, second] => (*hour, *minute, *second), + _ => return None, + }; + let (second, fraction) = match second.split_once('.') { + Some((second, fraction)) => { + if fraction.is_empty() || !fraction.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + (second, format!(".{fraction}")) + } + None => (second, String::new()), + }; + if ![hour, minute, second] + .iter() + .all(|part| !part.is_empty() && part.bytes().all(|b| b.is_ascii_digit())) + { + return None; + } + let hour = hour.parse().ok()?; + let minute = minute.parse().ok()?; + let second = second.parse().ok()?; + chrono::NaiveTime::from_hms_opt(hour, minute, second)?; + // dateTime 必须有 HH:MM:SS；补零表示该小时/分钟的起点，类型本身不能表达 + // 「只精确到小时/分钟」。小数秒不交给 chrono，避免超过纳秒的原文被截断。 + let lexical = format!("{calendar_date}T{hour:02}:{minute:02}:{second:02}{fraction}{zone}"); + Some(Literal::new_typed_literal(lexical, xsd::DATE_TIME)) +} + +/// parse_time 也接受 ±HH 和 ±HHMM，XSD 时区则要求 ±HH:MM。 +fn xsd_timezone(zone: &str) -> Option { + if !zone.is_ascii() { + return None; + } + let normalized = match zone.len() { + 3 => format!("{zone}:00"), + 5 => format!("{}:{}", &zone[..3], &zone[3..]), + 6 if zone.as_bytes()[3] == b':' => zone.to_string(), + _ => return None, }; - Some(Literal::new_typed_literal(text, ty)) + let digits = [&normalized[1..3], &normalized[4..6]]; + if !digits + .iter() + .all(|part| part.bytes().all(|b| b.is_ascii_digit())) + { + return None; + } + let hour: u32 = digits[0].parse().ok()?; + let minute: u32 = digits[1].parse().ok()?; + if hour > 14 || minute > 59 || (hour == 14 && minute != 0) { + return None; + } + Some(normalized) } /// 把事实的 `object_value` 形状抽出可写的字面文本，以及这段文本是不是人写的散文。 @@ -859,7 +973,7 @@ fn literal_text(v: &serde_json::Value) -> Option<(String, bool)> { #[cfg(test)] mod tests { use super::*; - use oxrdf::Quad; + use oxrdf::{NamedNodeRef, Quad}; fn kb() -> Uuid { Uuid::parse_str("01a06dc4-f40a-7013-b09f-1b499e2e7441").unwrap() @@ -945,7 +1059,7 @@ mod tests { } /// 导出一遍再解析回来。**必须解析回来**：断言字符串里有没有某一段， - /// 证明不了这份文件是不是合法的 Turtle，而那正是导出唯一要保证的事 + /// 证明不了文件语法是否合法；字面文本是否符合 XSD 类型仍须另行断言 fn export(format: Format, emit: impl FnOnce(&mut Sink, &Names, &Vocabulary)) -> Vec { let names = Names::new(kb(), None).unwrap(); let classes = vec![ @@ -1006,6 +1120,311 @@ mod tests { const OBJ: &str = ""; const WORKS_FOR: &str = "https://schema.org/worksFor"; + /// 同一个属性值穿过所有出口；RDF 能解析不代表 XSD 字面文本符合其类型。 + fn assert_attribute_literal_exports( + value: serde_json::Value, + datatype: &str, + expected: Option<(&str, NamedNodeRef<'_>)>, + ) { + let mut attribute = fact(5); + attribute.predicate_id = Some(id(4)); + attribute.object_id = None; + attribute.object_value = Some(value.clone()); + attribute + .qualifiers + .push(utopia_core::models::FactQualifier { + qualifier_type_id: id(4), + value: Some(value.clone()), + ..Default::default() + }); + let derived = ExportDerived { + id: id(7), + subject_id: id(10), + predicate_id: id(4), + object_id: None, + object_value: Some(value.clone()), + rule_id: None, + attribute_rule_id: Some(id(9)), + valid_from: None, + valid_from_precision: None, + valid_to: None, + valid_to_precision: None, + derived_at: at("2026-02-01T00:00:00Z"), + invalidated_at: None, + confidence: 0.9, + rule: "business".into(), + rule_predicate: None, + rule_name: Some("Attribute conclusion".into()), + premises: vec![id(5)], + premises_derived: Vec::new(), + subject_kb: Some(kb()), + object_kb: None, + predicate_kb: Some(kb()), + rule_kb: None, + attribute_rule_kb: Some(kb()), + foreign_fact_premise: false, + foreign_derived_premise: false, + subject_merged: false, + object_merged: false, + }; + let names = Names::new(kb(), None).unwrap(); + let mut property = relation(4, "value", None, "attribute"); + property.datatype = Some(datatype.into()); + let property_iri = names.relation(&property); + let statement_iri = names.fact(attribute.id).to_string(); + let subject_iri = names.entity(attribute.subject_id).to_string(); + let derived_iri = names.derived(derived.id).to_string(); + for format in [Format::Turtle, Format::JsonLd] { + let quads = export(format, |sink, names, _| { + let vocab = vocabulary(names, &[], std::slice::from_ref(&property)); + emit_fact(sink, names, &vocab, &attribute, at("2026-06-01T00:00:00Z")).unwrap(); + emit_derived(sink, names, &vocab, &derived).unwrap(); + }); + for (outlet, subject, predicate) in [ + ("fact", &statement_iri, rdf::OBJECT.as_str()), + ("current triple", &subject_iri, property_iri.as_str()), + ("qualifier", &statement_iri, property_iri.as_str()), + ("derived fact", &derived_iri, rdf::OBJECT.as_str()), + ] { + let objects: Vec<_> = quads + .iter() + .filter(|quad| { + quad.subject.to_string() == *subject && quad.predicate.as_str() == predicate + }) + .map(|quad| &quad.object) + .collect(); + let context = format!("{outlet}, {format:?}, {datatype}, {value}"); + match expected { + Some((lexical, expected_datatype)) => { + assert_eq!(objects.len(), 1, "{context}"); + let Term::Literal(literal) = objects[0] else { + panic!("expected a literal: {context}"); + }; + assert_eq!(literal.value(), lexical, "{context}"); + assert_eq!(literal.datatype(), expected_datatype, "{context}"); + } + None => assert!(objects.is_empty(), "unexpected object: {context}"), + } + } + } + } + + #[test] + fn normalized_numbers_export_as_decimal_without_exponents() { + for (input, expected) in [ + (serde_json::json!(0.0000001), "0.0000001"), + (serde_json::json!("1e-7"), "0.0000001"), + (serde_json::json!(1e20), "100000000000000000000"), + (serde_json::json!(-1e-7), "-0.0000001"), + (serde_json::json!(-1e20), "-100000000000000000000"), + (serde_json::json!(0), "0.0"), + (serde_json::json!(-0.0), "-0.0"), + (serde_json::json!(12.375), "12.375"), + (serde_json::json!(-12.375), "-12.375"), + (serde_json::json!(1.25e-7), "0.000000125"), + ] { + let normalized = utopia_extract::normalize_attr_value("number", &input).unwrap(); + assert_attribute_literal_exports( + normalized.clone(), + "number", + Some((expected, xsd::DECIMAL)), + ); + assert_attribute_literal_exports( + serde_json::json!({"value": normalized}), + "number", + Some((expected, xsd::DECIMAL)), + ); + } + // 已存的整数不能再经过 f64：这个值超过了双精度能逐个表示整数的范围。 + assert_attribute_literal_exports( + serde_json::json!({"value": 9_007_199_254_740_993_u64}), + "number", + Some(("9007199254740993", xsd::DECIMAL)), + ); + // 最小次正规数和最大有限数：不能因固定小数位变成零，也不能溢出。 + for (input, expected) in [ + (f64::from_bits(1), format!("0.{}5", "0".repeat(323))), + (f64::MAX, format!("17976931348623157{}", "0".repeat(292))), + ] { + let normalized = + utopia_extract::normalize_attr_value("number", &serde_json::json!(input)).unwrap(); + assert_attribute_literal_exports(normalized, "number", Some((&expected, xsd::DECIMAL))); + } + } + + #[test] + fn normalized_dates_export_their_precision_and_original_clock() { + for (input, expected, datatype) in [ + ("2024", "2024", xsd::G_YEAR), + ("2024-07", "2024-07", xsd::G_YEAR_MONTH), + ("2024-07-15", "2024-07-15", xsd::DATE), + ( + "2024-07-15T14:32:07Z", + "2024-07-15T14:32:07Z", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14:32:07+08:00", + "2024-07-15T14:32:07+08:00", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14:32:07.382-05:30", + "2024-07-15T14:32:07.382-05:30", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14:32:07.123456789012Z", + "2024-07-15T14:32:07.123456789012Z", + xsd::DATE_TIME, + ), + ("2024-07-15T14Z", "2024-07-15T14:00:00Z", xsd::DATE_TIME), + ( + "2024-07-15T14+08:00", + "2024-07-15T14:00:00+08:00", + xsd::DATE_TIME, + ), + ("2024-07-15T14:32Z", "2024-07-15T14:32:00Z", xsd::DATE_TIME), + ( + "2024-07-15T14:32+08:00", + "2024-07-15T14:32:00+08:00", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14:32+0800", + "2024-07-15T14:32:00+08:00", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14+08", + "2024-07-15T14:00:00+08:00", + xsd::DATE_TIME, + ), + ("2024-07-15T14:32z", "2024-07-15T14:32:00Z", xsd::DATE_TIME), + ( + "2024-07-15T14+14", + "2024-07-15T14:00:00+14:00", + xsd::DATE_TIME, + ), + ( + "2024-07-15T14-14", + "2024-07-15T14:00:00-14:00", + xsd::DATE_TIME, + ), + ( + "2024-07-15 14:32+08:00", + "2024-07-15T14:32:00+08:00", + xsd::DATE_TIME, + ), + ("2024-07-15T14", "2024-07-15T14:00:00", xsd::DATE_TIME), + ("2024-07-15T14:32", "2024-07-15T14:32:00", xsd::DATE_TIME), + ( + "2024-07-15T14:32:07.382", + "2024-07-15T14:32:07.382", + xsd::DATE_TIME, + ), + ] { + let normalized = + utopia_extract::normalize_attr_value("date", &serde_json::json!(input)).unwrap(); + assert_attribute_literal_exports( + normalized.clone(), + "date", + Some((expected, datatype)), + ); + assert_attribute_literal_exports( + serde_json::json!({"value": normalized}), + "date", + Some((expected, datatype)), + ); + } + } + + #[test] + fn attribute_calendar_years_are_not_cut_to_four_characters() { + for (input, expected, datatype) in [ + ("-0001-07", "-0001-07", xsd::G_YEAR_MONTH), + ("-0001-07-15", "-0001-07-15", xsd::DATE), + ("+10000-07", "10000-07", xsd::G_YEAR_MONTH), + ("+10000-07-15", "10000-07-15", xsd::DATE), + ( + "+10000-07-15T14:32:07.382Z", + "10000-07-15T14:32:07.382Z", + xsd::DATE_TIME, + ), + ] { + let normalized = + utopia_extract::normalize_attr_value("date", &serde_json::json!(input)).unwrap(); + assert_attribute_literal_exports(normalized, "date", Some((expected, datatype))); + } + } + + #[test] + fn unrecognized_literal_text_keeps_the_existing_fallback() { + for value in ["1e309", "1e-325", "1e2147483648", "1_0e2", "before"] { + assert_eq!(decimal_text(value.to_string()), value); + } + for value in [ + "2024-07-15Tnot-a-clock", + "2024-07-15T14:32:07.noiseZ", + "2024-07-15T14+14:01", + "2024-07-15T14+08:60", + "2024-07-15T14+时区", + ] { + assert!(date_literal(value).is_none(), "{value}"); + } + } + + #[test] + fn typed_attributes_keep_value_precedence_prose_and_missing_objects() { + for (value, datatype, expected) in [ + ( + serde_json::json!({"value": 12.5, "summary": "estimate"}), + "number", + Some(("12.5", xsd::DECIMAL)), + ), + ( + serde_json::json!({"value": "2024", "summary": "sometime in 2024"}), + "date", + Some(("2024", xsd::G_YEAR)), + ), + ( + serde_json::json!({"value": 0, "summary": ""}), + "number", + Some(("0", xsd::DECIMAL)), + ), + ( + serde_json::json!({"summary": "before the merge"}), + "date", + Some(("before the merge", xsd::STRING)), + ), + ( + serde_json::json!({"value": null, "summary": "not specified"}), + "number", + Some(("not specified", xsd::STRING)), + ), + ( + serde_json::json!({"class": "gas_well"}), + "number", + Some(("gas_well", xsd::STRING)), + ), + ( + serde_json::json!({"value": "45 days after the Trigger Date", "relative": true}), + "date", + Some(("45 days after the Trigger Date", xsd::STRING)), + ), + (serde_json::json!({"value": null}), "date", None), + (serde_json::json!({"confidence": 0.9}), "number", None), + ( + serde_json::json!({"summary": "", "class": ""}), + "date", + None, + ), + (serde_json::Value::Null, "number", None), + ] { + assert_attribute_literal_exports(value, datatype, expected); + } + } + #[test] fn unbound_literal_objects_survive_both_formats() { for value in [ @@ -1118,13 +1537,13 @@ mod tests { let cases: &[(&str, serde_json::Value, &[&str])] = &[ ( "value wins over summary when both are present", - serde_json::json!({ "value": "2026-01-15", "summary": "around mid-January" }), - &["\"2026-01-15\"^^"], + serde_json::json!({ "value": 12.5, "summary": "about a dozen" }), + &["\"12.5\"^^"], ), ( "value alone resolves to its scalar string", - serde_json::json!({ "value": "45 days after the Trigger Date" }), - &["\"45 days after the Trigger Date\"^^"], + serde_json::json!({ "value": "42" }), + &["\"42\"^^"], ), ( "summary alone is prose: a string, never the declared type", @@ -1133,8 +1552,8 @@ mod tests { ), ( "empty summary with a value resolves to the value", - serde_json::json!({ "summary": "", "value": "actual text" }), - &["\"actual text\"^^"], + serde_json::json!({ "summary": "", "value": "0" }), + &["\"0\"^^"], ), ( "null value and no summary → no rdf:object", diff --git a/crates/utopia-server/src/time_resolution.rs b/crates/utopia-server/src/time_resolution.rs index ece296b19..c7f386b89 100644 --- a/crates/utopia-server/src/time_resolution.rs +++ b/crates/utopia-server/src/time_resolution.rs @@ -13,12 +13,16 @@ //! 没有一个时间词写在代码里（决定 8）。写明的日期是 A 级，从文档自己给的锚点算出来的是 //! B 级，锚不到的是 C 级——C 级什么也不写，等锚点（决定 4、6）。 //! 4. 结果落到开放陈述的 `valid_from` / `valid_to`（决定 5 的第一根轴），时间轴上才有它。 +//! +//! 例外只长在推送的陈述上（0054 来源）：`when` / `ended` 写的是完整日期或带区时刻时， +//! 载荷就已经是结构化的绝对值——结构化的时间戳不是时间词（0045 决定 8，#1089），代码 +//! 直接读成解释，不送模型；没有对话模型时，代码读得动的照解，读不动的留 C，任务照常收工。 use crate::extraction::chat_retrying_rate_limits_at; use crate::llm_util; use crate::state::AppState; use chrono::{DateTime, Duration, Months, NaiveDate, NaiveTime, TimeZone, Utc}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use utopia_extract::time::{ build_interpretation_messages, headings_at, now_in_force, parse_interpretation_response, Anchor, DateParts, Direction, DocumentDating, Granularity, Interpretation, MentionInput, @@ -232,7 +236,7 @@ fn anchor_time( } } -/// 一条解释 → 世界轴位置。`earlier` 是同一批里已经算出来的提及（锚点指着它们）。 +/// 一条解释 → 世界轴位置。`earlier` 是同一文档里已经算出来的提及（锚点指着它们）。 fn resolve_one( interp: &Interpretation, ctx: &DocumentDating, @@ -347,6 +351,89 @@ fn resolve_one( } } +/// 计算已收到的解释；每条解释使用它自己的章节上下文。 +fn resolve_interpretations( + interpretations: &[Interpretation], + context_for: impl Fn(i64) -> DocumentDating, +) -> HashMap { + let mut resolved = HashMap::new(); + let mut pending: Vec<&Interpretation> = interpretations.iter().collect(); + while !pending.is_empty() { + let before = pending.len(); + pending.retain(|interp| { + let r = resolve_one(interp, &context_for(interp.id), &resolved); + if r == UNRESOLVED { + return true; + } + resolved.insert(interp.id, r); + false + }); + // 每轮至少解开一条才继续；循环、自引用、缺失或不可用的锚点都在无进展时留下 C。 + // 不能固定遍数：模型回复可以把任意长的依赖链倒着写。 + if pending.len() == before { + break; + } + } + for interp in pending { + resolved.insert(interp.id, UNRESOLVED); + } + resolved +} + +/// 算的时候文档的日期换成提及所在那一节的起算点。 +fn context_at(context: &DocumentDating, now: Option) -> DocumentDating { + let mut scoped = context.clone(); + if let Some(entry) = now.and_then(|i| context.entries.get(i)) { + scoped.date = Some(entry.from.clone()); + scoped.date_words = Some(entry.words.clone()); + } + scoped +} + +/// 推送的陈述把 `when` / `ended` 写成完整日期或带区时刻时，这里就把它读成一条解释 +/// （#1089；0045 决定 8：结构化的时间戳不是时间词）。形状是点、参照是照写的绝对值、 +/// 粒度是字写到的那一级——与模型对同样的字该给的回答一致。`from` 存换算成 UTC 后的 +/// 部件：`DateParts` 不带区，偏移在 `parse_time` 里已经折进值里；存 UTC 部件的解释 +/// 再算一遍得到同一个值，存原样的本地钟点反而会读出另一个时刻。 +/// +/// 只收 `parse_time` 读到「日」或更细的：`YYYY-MM-DD` 与带 `Z` 或偏移的 +/// `YYYY-MM-DDTHH[:MM[:SS]]`。光杆的年、年月、以及任何读不动的字都回 `None`， +/// 照旧走模型那条路——这一刀不替代码认时间词。 +fn stated_reading(text: &str, id: i64) -> Option { + let (t, precision) = utopia_extract::parse_time(text)?; + use chrono::{Datelike, Timelike}; + let (granularity, hour, minute, second) = match precision { + "day" => (Granularity::Day, None, None, None), + "hour" => (Granularity::Hour, Some(t.hour()), None, None), + "minute" => (Granularity::Minute, Some(t.hour()), Some(t.minute()), None), + "second" => ( + Granularity::Second, + Some(t.hour()), + Some(t.minute()), + Some(t.second()), + ), + // 光杆的年、年月不在这刀里（#1089 的界），退回 None 让模型读 + _ => return None, + }; + Some(Interpretation { + id, + shape: Shape::Point, + reference: Reference::Absolute { + from: DateParts { + year: t.year(), + quarter: None, + month: Some(t.month()), + day: Some(t.day()), + hour, + minute, + second, + }, + to: None, + }, + granularity, + }) +} + /// 一条陈述的起与止：`when` 给起（点与区间也给止），`ended` 给止。 fn combine(when: Option, ended: Option) -> Resolved { let mut r = when.unwrap_or(UNRESOLVED); @@ -373,11 +460,17 @@ pub async fn resolve_document(state: &AppState, document_id: Uuid) -> anyhow::Re return Ok(()); } let kb = utopia_store::kbs::get(pool, doc.kb_id).await?; - let settings = utopia_store::settings::get(pool, kb.workspace_id) + let settings = utopia_store::settings::get(pool, kb.workspace_id).await?; + let client = settings.as_ref().and_then(llm_util::chat_client); + // 推送来的陈述（0054）：`when` / `ended` 写的完整日期或带区时刻由代码直接读（#1089）， + // 没有对话模型也照常解；别的文档照旧——没有对话模型，这一步就没有替它读字的人 + let pushed = crate::pipeline::source_kind(state, doc.source_id) .await? - .ok_or_else(|| anyhow::anyhow!("Chat model not configured; cannot resolve time"))?; - let client = llm_util::chat_client(&settings) - .ok_or_else(|| anyhow::anyhow!("Chat model not configured; cannot resolve time"))?; + .as_deref() + == Some("statements"); + if !pushed && client.is_none() { + anyhow::bail!("Chat model not configured; cannot resolve time"); + } // 1. 文档时间上下文：抽取时各块报上来的（0064）。这里不问模型 let context: DocumentDating = match doc.time_context.clone() { @@ -465,109 +558,107 @@ pub async fn resolve_document(state: &AppState, document_id: Uuid) -> anyhow::Re let mut distinct: Vec = Vec::new(); let mut index: HashMap = HashMap::new(); let mut key_of: HashMap = HashMap::new(); + // 推送来的陈述（0054）：载荷写的完整日期或带区时刻由代码读成解释（#1089）。 + // 解释与别的提及同去同回，只是它从不进问模型的清单——重问也只能是这个答案， + // 代码读过的不许模型盖掉；判据是文档来自 `statements` 来源，不是字长得像日期， + // 普通文档里同样的字照旧交给模型 + let mut presolved: HashMap = HashMap::new(); for m in &mentions { let key = (m.text.clone(), m.sentence.clone(), now_of(m)); key_of.insert(m.id, key.clone()); - if !index.contains_key(&key) { - index.insert(key.clone(), distinct.len()); - distinct.push(key); + if index.contains_key(&key) { + continue; } - } - // 算的时候文档的日期换成那一节的起算点 - let context_at = |now: Option| -> DocumentDating { - let mut c = context.clone(); - if let Some(e) = now.and_then(|i| context.entries.get(i)) { - c.date = Some(e.from.clone()); - c.date_words = Some(e.words.clone()); + index.insert(key.clone(), distinct.len()); + distinct.push(key); + if pushed { + if let Some(interp) = stated_reading(&m.text, (distinct.len() - 1) as i64) { + presolved.insert(distinct.len() - 1, interp); + } } - c - }; + } let ctx = TimeContext { date: context.date.as_ref(), date_words: context.date_words.as_deref(), periods: &context.periods, fiscal_year_end: context.fiscal_year_end, }; - let mut resolved: HashMap = HashMap::new(); + let mut interpretations: Vec = presolved.values().cloned().collect(); + let mut received = HashSet::new(); let mut malformed = 0usize; // 问两轮：第一轮全部，第二轮只问第一轮没答到的。模型漏答的提及从前就留着字、没有解释， - // 也没有人再问（测量台上的「上周」三遍都是这样） - let mut todo: Vec = (0..distinct.len()).collect(); - for round in 0..2 { - if round == 1 { - todo.retain(|i| !resolved.contains_key(i)); - if todo.is_empty() { - break; - } - tracing::info!(%document_id, unanswered = todo.len(), "时间解释有没答到的，再问一次"); - } - for batch in todo.clone().chunks(BATCH) { - let inputs: Vec> = batch - .iter() - .map(|&i| MentionInput { - id: i as i64, - text: &distinct[i].0, - sentence: &distinct[i].1, - }) - .collect(); - let ids: Vec = inputs.iter().map(|m| m.id).collect(); - let messages = build_interpretation_messages(&ctx, &inputs); - let reply = match chat_retrying_rate_limits_at( - state, - &settings, - &client, - &messages, - Some(0.0), - ) - .await - { - Ok(r) => r, - Err(e) => { - tracing::warn!(%document_id, error = %e, "时间解释调用失败，这一批留作未解析"); - continue; + // 也没有人再问（测量台上的「上周」三遍都是这样）。代码已经读出答案的不问 + let mut todo: Vec = (0..distinct.len()) + .filter(|i| !presolved.contains_key(i)) + .collect(); + if let (Some(settings), Some(client)) = (settings.as_ref(), client.as_ref()) { + for round in 0..2 { + if round == 1 { + todo.retain(|i| !received.contains(&(*i as i64))); + if todo.is_empty() { + break; } - }; - let (interps, skipped) = match parse_interpretation_response(&reply.text, &ids) { - Ok(x) => x, - Err(e) => { - tracing::warn!(%document_id, error = %e, "时间解释回复解析失败，这一批留作未解析"); - continue; - } - }; - malformed += skipped; - if skipped > 0 { - // 坏条目要看得见是哪一条、坏在哪：只报一个数的话，「上周」三遍都没日期也查不出原因 - tracing::warn!(%document_id, skipped, reply = %reply.text.chars().take(600).collect::(), "时间解释里有读不了的条目"); + tracing::info!(%document_id, unanswered = todo.len(), "时间解释有没答到的，再问一次"); } - // 两遍：先算不靠别的提及的，再算指着别的提及的 - let mut earlier: HashMap = HashMap::new(); - let anchored_to_mention = |i: &Interpretation| { - matches!( - &i.reference, - Reference::Anchored { - anchor: Anchor::Mention { .. }, - .. - } - ) - }; - for pass in 0..2 { - for interp in interps + for batch in todo.clone().chunks(BATCH) { + let inputs: Vec> = batch .iter() - .filter(|i| anchored_to_mention(i) == (pass == 1)) + .map(|&i| MentionInput { + id: i as i64, + text: &distinct[i].0, + sentence: &distinct[i].1, + }) + .collect(); + let ids: Vec = inputs.iter().map(|m| m.id).collect(); + let messages = build_interpretation_messages(&ctx, &inputs); + let reply = match chat_retrying_rate_limits_at( + state, + settings, + client, + &messages, + Some(0.0), + ) + .await { - let now = usize::try_from(interp.id) - .ok() - .and_then(|i| distinct.get(i)) - .and_then(|k| k.2); - let r = resolve_one(interp, &context_at(now), &earlier); - earlier.insert(interp.id, r); - if let Ok(i) = usize::try_from(interp.id) { - resolved.insert(i, (interp.clone(), r)); + Ok(r) => r, + Err(e) => { + tracing::warn!(%document_id, error = %e, "时间解释调用失败，这一批留作未解析"); + continue; + } + }; + let (interps, skipped) = match parse_interpretation_response(&reply.text, &ids) { + Ok(x) => x, + Err(e) => { + tracing::warn!(%document_id, error = %e, "时间解释回复解析失败，这一批留作未解析"); + continue; } + }; + malformed += skipped; + if skipped > 0 { + // 坏条目要看得见是哪一条、坏在哪：只报一个数的话，「上周」三遍都没日期也查不出原因 + tracing::warn!(%document_id, skipped, reply = %reply.text.chars().take(600).collect::(), "时间解释里有读不了的条目"); } + received.extend(interps.iter().map(|interp| interp.id)); + interpretations.extend(interps); } } } + // 收到解释和算出日期是两件事：补答只问没收到的，计算则看全部批次和补答里的解释。 + // 这样下游不用再问模型，也能用上后一批或补答才给出的锚点。 + let computed = resolve_interpretations(&interpretations, |id| { + let now = usize::try_from(id) + .ok() + .and_then(|i| distinct.get(i)) + .and_then(|key| key.2); + context_at(&context, now) + }); + let mut resolved: HashMap = HashMap::new(); + for interp in interpretations { + if let Ok(i) = usize::try_from(interp.id) { + let r = computed[&interp.id]; + resolved.insert(i, (interp, r)); + } + } // 3. 写回提及，再按陈述合起止 let mut per_fact: HashMap, Option)> = HashMap::new(); @@ -577,7 +668,23 @@ pub async fn resolve_document(state: &AppState, document_id: Uuid) -> anyhow::Re continue; }; let Some((interp, r)) = resolved.get(&i) else { - // 这一批没答到：留着字，不写等级（下次重跑再问） + // 推送的文档没有对话模型可问（#1089）：代码读不动的提及按 C 记——「有提及、 + // 没能放上轴」和「没有时间词」要分得出；任务照常收工，不因一个没配的模型 + // 留下永远失败的任务。解释列留空：没有人读过它，不编一条出来 + if pushed && client.is_none() { + utopia_store::time_mentions::set_resolution( + pool, m.id, "C", None, None, None, None, + ) + .await?; + c += 1; + let slot = per_fact.entry(m.fact_id).or_insert((None, None)); + if m.role == "ended" { + slot.1 = Some(UNRESOLVED); + } else { + slot.0 = Some(UNRESOLVED); + } + } + // 有模型而没答到的：留着字，不写等级（下次重跑再问） continue; }; utopia_store::time_mentions::set_interpretation( @@ -642,6 +749,17 @@ pub async fn resolve_document(state: &AppState, document_id: Uuid) -> anyhow::Re Ok(()) } +/// 见证的名字：日期的名字加它的字（「提报日期 2026年9月4日」）。没有名字的光是字；封面上 +/// 一个光秃秃的日期，模型常把日期自己又填成名字，那就只写一遍 +fn witness_name(name: &str, words: &str) -> String { + let (name, words) = (name.trim(), words.trim()); + if name.is_empty() || name == words { + words.to_string() + } else { + format!("{name} {words}") + } +} + /// 这篇文档的陈述各由它所在那一节里文本说话的那一刻作证（0064 决定 3）。 /// /// 一条陈述在哪一节，看它的引文在它那一块里的位置、那一块正文里的标题；管着它的是范围 @@ -680,12 +798,7 @@ async fn attest_statements( let path = headings_at(&text, usize::try_from(start.unwrap_or(0)).unwrap_or(0)); let witness = now_in_force(&context.entries, &path).and_then(|now| { let (at, _) = parts_to_time(&now.from)?; - let by = if now.name.is_empty() { - now.words.clone() - } else { - format!("{} {}", now.name, now.words) - }; - Some((at, by)) + Some((at, witness_name(&now.name, &now.words))) }); // 每条证据记下自己这一节的日期（没说的记成空）：事实上只留最早的那一个，证据所在 // 的文档删了之后要照还在的证据重算，得知道每一条各是哪天 @@ -727,9 +840,84 @@ mod context_tests; #[path = "time_attestation_tests.rs"] mod attestation_tests; +#[cfg(test)] +#[path = "time_resolution_pushed_tests.rs"] +mod pushed_tests; + #[cfg(test)] mod tests { use super::*; + + /// 推送载荷里写明的值（#1089）：完整日期与带区时刻读成点状绝对解释，精度随字； + /// 光杆的年、年月、读不动的字都回 None——不替代码认时间词 + #[test] + fn a_stated_value_reads_at_the_precision_it_writes() { + let read = |s: &str| { + let interp = stated_reading(s, 0)?; + let Reference::Absolute { from, to } = interp.reference else { + return None; + }; + assert_eq!(interp.shape, Shape::Point); + assert!(to.is_none()); + Some((from, interp.granularity)) + }; + let date = |y, m, d, h, mi, s| DateParts { + year: y, + quarter: None, + month: Some(m), + day: Some(d), + hour: h, + minute: mi, + second: s, + }; + assert_eq!( + read("2026-09-23"), + Some((date(2026, 9, 23, None, None, None), Granularity::Day)) + ); + assert_eq!( + read("2026-09-23T08Z"), + Some((date(2026, 9, 23, Some(8), None, None), Granularity::Hour)) + ); + assert_eq!( + read("2026-09-23T08:14Z"), + Some(( + date(2026, 9, 23, Some(8), Some(14), None), + Granularity::Minute + )) + ); + assert_eq!( + read("2026-09-23T08:14:03Z"), + Some(( + date(2026, 9, 23, Some(8), Some(14), Some(3)), + Granularity::Second + )) + ); + // 偏移折进值里：+08:00 的 08:14 是 UTC 的 00:14 + assert_eq!( + read("2026-09-23T08:14:03+08:00"), + Some(( + date(2026, 9, 23, Some(0), Some(14), Some(3)), + Granularity::Second + )) + ); + // 这一刀之外的字：光杆的年、年月、自然语言 + for s in ["2026", "2026-09", "last Tuesday", "ended", ""] { + assert!(read(s).is_none(), "{s:?} is not code-resolved"); + } + } + + #[test] + fn a_bare_date_is_named_once() { + assert_eq!( + witness_name("提报日期", "2026年9月4日"), + "提报日期 2026年9月4日" + ); + assert_eq!(witness_name("", "2024年10月10日"), "2024年10月10日"); + assert_eq!( + witness_name("2024年10月10日", " 2024年10月10日"), + "2024年10月10日" + ); + } use utopia_extract::time::NamedPeriod; fn parts(y: i32, m: Option, d: Option) -> DateParts { @@ -761,6 +949,268 @@ mod tests { } } + fn day_point(date: &str, grade: &'static str) -> Resolved { + Resolved { + from: Some(at(date)), + from_p: Some("day"), + to: Some(at(date)), + to_p: Some("day"), + grade, + } + } + + fn absolute_day(id: i64) -> Interpretation { + Interpretation { + id, + shape: Shape::Point, + reference: Reference::Absolute { + from: parts(2024, Some(1), Some(1)), + to: None, + }, + granularity: Granularity::Day, + } + } + + fn next_day(id: i64, anchor_id: i64) -> Interpretation { + Interpretation { + reference: Reference::Anchored { + anchor: Anchor::Mention { id: anchor_id }, + offset: Some(Offset { + count: 1, + unit: Unit::Day, + direction: Direction::After, + }), + }, + ..absolute_day(id) + } + } + + #[test] + fn mention_dependencies_resolve_in_all_reply_orders() { + let items = [ + serde_json::json!([0, "point", {"kind": "absolute", "from": {"y": 2024, "m": 1, "d": 1}}, "day"]), + serde_json::json!([1, "point", {"kind": "anchored", "anchor": {"kind": "mention", "id": 0}, "offset": {"count": 1, "unit": "day", "direction": "after"}}, "day"]), + serde_json::json!([2, "point", {"kind": "anchored", "anchor": {"kind": "mention", "id": 1}, "offset": {"count": 1, "unit": "day", "direction": "after"}}, "day"]), + ]; + let expected = HashMap::from([ + (0, day_point("2024-01-01T00:00:00Z", "A")), + (1, day_point("2024-01-02T00:00:00Z", "B")), + (2, day_point("2024-01-03T00:00:00Z", "B")), + ]); + for order in [ + [0, 1, 2], + [0, 2, 1], + [1, 0, 2], + [1, 2, 0], + [2, 0, 1], + [2, 1, 0], + ] { + let reply = serde_json::json!({"m": order.map(|i| items[i].clone())}); + let (interpretations, skipped) = + parse_interpretation_response(&reply.to_string(), &[0, 1, 2]).unwrap(); + assert_eq!(skipped, 0); + assert_eq!(interpretations.len(), 3); + let resolved = resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!(resolved, expected, "reply order: {order:?}"); + } + } + + #[test] + fn a_long_reverse_chain_does_not_follow_numeric_ids_or_a_fixed_pass_limit() { + // 编号上下交错，与依赖次序无关；链长也超过一批的大小。 + let ids: Vec = (0..81).map(|depth| 500 + (depth * 37) % 101).collect(); + let mut interpretations = vec![absolute_day(ids[0])]; + for pair in ids.windows(2) { + interpretations.push(next_day(pair[1], pair[0])); + } + interpretations.reverse(); + let resolved = resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!(resolved.len(), ids.len()); + for (depth, id) in ids.iter().enumerate() { + let date = at("2024-01-01T00:00:00Z") + Duration::days(depth as i64); + assert_eq!( + resolved[id], + Resolved { + from: Some(date), + from_p: Some("day"), + to: Some(date), + to_p: Some("day"), + grade: if depth == 0 { "A" } else { "B" }, + }, + "depth {depth}, id {id}" + ); + } + } + + #[test] + fn independent_chains_resolve_while_cycles_and_missing_anchors_stay_unresolved() { + let second_root = Interpretation { + reference: Reference::Absolute { + from: parts(2025, Some(6), Some(10)), + to: None, + }, + ..absolute_day(20) + }; + let interpretations = [ + next_day(2, 1), + next_day(30, 30), // 自引用 + next_day(31, 32), + next_day(32, 31), // 循环 + next_day(33, 999), // 缺失的解释 + next_day(34, 31), // 依赖循环的下游 + next_day(22, 21), + next_day(1, 0), + next_day(21, 20), + second_root, + absolute_day(0), + ]; + let resolved = resolve_interpretations(&interpretations, |_| ctx_dated(2030, 1, 1)); + assert_eq!( + resolved, + HashMap::from([ + (0, day_point("2024-01-01T00:00:00Z", "A")), + (1, day_point("2024-01-02T00:00:00Z", "B")), + (2, day_point("2024-01-03T00:00:00Z", "B")), + (20, day_point("2025-06-10T00:00:00Z", "A")), + (21, day_point("2025-06-11T00:00:00Z", "B")), + (22, day_point("2025-06-12T00:00:00Z", "B")), + (30, UNRESOLVED), + (31, UNRESOLVED), + (32, UNRESOLVED), + (33, UNRESOLVED), + (34, UNRESOLVED), + ]) + ); + } + + #[test] + fn received_interpretations_can_use_anchors_from_later_batches_and_the_makeup_reply() { + let first_reply = r#"{"m": [ + [2, "point", {"kind": "anchored", "anchor": {"kind": "mention", "id": 1}, "offset": {"count": 1, "unit": "day", "direction": "after"}}, "day"], + [1, "point", {"kind": "anchored", "anchor": {"kind": "mention", "id": 40}, "offset": {"count": 1, "unit": "day", "direction": "after"}}, "day"] + ]}"#; + let (mut interpretations, skipped) = + parse_interpretation_response(first_reply, &[0, 1, 2]).unwrap(); + assert_eq!(skipped, 0); + let before = resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!(before, HashMap::from([(1, UNRESOLVED), (2, UNRESOLVED)])); + assert!(!before.contains_key(&0), "未回答与已回答但未解析的提及不同"); + + let later_batch = r#"{"m": [ + [40, "point", {"kind": "anchored", "anchor": {"kind": "mention", "id": 0}, "offset": {"count": 1, "unit": "day", "direction": "after"}}, "day"] + ]}"#; + let (later, skipped) = parse_interpretation_response(later_batch, &[40]).unwrap(); + assert_eq!(skipped, 0); + interpretations.extend(later); + let still_waiting = + resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!(still_waiting[&2], UNRESOLVED); + + let makeup_reply = r#"{"m": [ + [0, "point", {"kind": "absolute", "from": {"y": 2024, "m": 1, "d": 1}}, "day"] + ]}"#; + let (makeup, skipped) = parse_interpretation_response(makeup_reply, &[0]).unwrap(); + assert_eq!(skipped, 0); + interpretations.extend(makeup); + let after = resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!( + after, + HashMap::from([ + (0, day_point("2024-01-01T00:00:00Z", "A")), + (40, day_point("2024-01-02T00:00:00Z", "B")), + (1, day_point("2024-01-03T00:00:00Z", "B")), + (2, day_point("2024-01-04T00:00:00Z", "B")), + ]) + ); + } + + #[test] + fn dependency_resolution_keeps_each_mentions_section_context() { + let mut context = ctx_dated(2030, 1, 1); + for (section, year) in [("First report", 2024), ("Second report", 2025)] { + context.entries.push(utopia_extract::time::TimeEntry { + kind: "now".into(), + name: String::new(), + words: year.to_string(), + from: parts(year, Some(1), Some(1)), + to: None, + scope: vec![section.into()], + chunk: None, + char_start: None, + }); + } + let document_anchor = |id| Interpretation { + reference: Reference::Anchored { + anchor: Anchor::Document, + offset: None, + }, + ..absolute_day(id) + }; + // 偶数在第一节，奇数在第二节；两条链各自从本节日期起算。 + let interpretations = [ + next_day(4, 2), + next_day(5, 3), + next_day(2, 0), + next_day(3, 1), + document_anchor(1), + document_anchor(0), + ]; + let resolved = resolve_interpretations(&interpretations, |id| { + context_at(&context, Some((id % 2) as usize)) + }); + assert_eq!( + resolved, + HashMap::from([ + (0, day_point("2024-01-01T00:00:00Z", "B")), + (2, day_point("2024-01-02T00:00:00Z", "B")), + (4, day_point("2024-01-03T00:00:00Z", "B")), + (1, day_point("2025-01-01T00:00:00Z", "B")), + (3, day_point("2025-01-02T00:00:00Z", "B")), + (5, day_point("2025-01-03T00:00:00Z", "B")), + ]) + ); + } + + #[test] + fn dependency_resolution_preserves_dateless_endings_and_unusable_anchors() { + let interpretations = [ + next_day(1, 0), + Interpretation { + shape: Shape::Until, + ..absolute_day(0) + }, + next_day(3, 2), + Interpretation { + shape: Shape::EndedUnknown, + ..next_day(2, 2) + }, + ]; + let resolved = resolve_interpretations(&interpretations, |_| DocumentDating::default()); + assert_eq!( + resolved, + HashMap::from([ + ( + 0, + Resolved { + from: None, + from_p: None, + ..day_point("2024-01-01T00:00:00Z", "A") + } + ), + (1, UNRESOLVED), + ( + 2, + Resolved { + to_p: Some(ENDED_UNKNOWN), + grade: "A", + ..UNRESOLVED + } + ), + (3, UNRESOLVED), + ]) + ); + } + #[test] fn parts_give_the_precision_the_words_state() { assert_eq!(parts_to_time(&parts(2024, None, None)).unwrap().1, "year"); diff --git a/crates/utopia-server/src/time_resolution_pushed_tests.rs b/crates/utopia-server/src/time_resolution_pushed_tests.rs new file mode 100644 index 000000000..02d3d1062 --- /dev/null +++ b/crates/utopia-server/src/time_resolution_pushed_tests.rs @@ -0,0 +1,600 @@ +//! 推送的陈述把 `when` / `ended` 写成完整日期或带区时刻时，`resolve_time` 自己读（#1089， +//! 0054 的悬而未决那条的落点）：不问模型，没有对话模型也照常收工。 +//! +//! 与 `time_context_tests.rs` 同一个做法：脚本化的模型端点、真库， +//! 没有 `UTOPIA_DATABASE_URL` 时跳过而不是失败。自建自拆，绝不碰已有的库。 +use super::*; +use axum::{extract::State, response::IntoResponse, routing::post, Json, Router}; +use serde_json::{json, Value}; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +struct Model { + replies: Arc>>, + requests: Arc>>, +} + +async fn reply(State(m): State, Json(body): Json) -> impl IntoResponse { + m.requests.lock().unwrap().push(body); + let text = { + let mut replies = m.replies.lock().unwrap(); + if replies.is_empty() { + panic!("unexpected model request"); + } + replies.remove(0) + }; + let frame = json!({"choices":[{"delta":{"content":text}}]}); + let done = json!({"choices":[{"delta":{},"finish_reason":"stop"}]}); + ( + [("content-type", "text/event-stream")], + format!("data: {frame}\n\ndata: {done}\n\ndata: [DONE]\n\n"), + ) +} + +/// 陈述的开放行：(宾语名, valid_from, from 精度, valid_to, to 精度, 起点等级) +type FactRow = ( + String, + Option>, + Option, + Option>, + Option, + Option, +); + +/// 一条时间提及：(字, 槽, 等级, 形状) +type MentionRow = (String, String, Option, Option); + +/// 一个 `statements` 来源、一篇推送来的文档（块就是载荷）、随需一个脚本化的对话端点。 +/// `replies` 为 None 时不配模型——这正是这条路的测试点 +struct Rig { + pool: sqlx::PgPool, + state: AppState, + org: Uuid, + kb: Uuid, + doc: Uuid, + model: Option, + server: Option>, + _dir: tempfile::TempDir, +} + +impl Rig { + async fn new(payload: Value, replies: Option>) -> anyhow::Result> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(None); + }; + let pool = sqlx::PgPool::connect(&url).await?; + utopia_store::db::migrate(&pool).await?; + let (org, ws, kb, source, doc, chunk) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + // 存的文档是 `push_statements` 自己重排出来的那份：身份、日期（有的话）， + // 再是契约的三个数组，键序固定 + let text = serde_json::to_string_pretty(&payload)?; + sqlx::raw_sql(&format!( + "INSERT INTO organizations(id,name) VALUES ('{org}','pushed-time'); + INSERT INTO workspaces(id,org_id,name) VALUES ('{ws}','{org}','pushed-time'); + INSERT INTO knowledge_bases(id,workspace_id,name) VALUES ('{kb}','{ws}','pushed-time'); + INSERT INTO sources(id,kb_id,kind,name) VALUES ('{source}','{kb}','statements','robot-1'); + INSERT INTO documents(id,kb_id,filename,sha256,source_id) VALUES + ('{doc}','{kb}','obs-1.json','x{doc}','{source}'); + INSERT INTO chunks(id,kb_id,document_id,seq,text) VALUES + ('{chunk}','{kb}','{doc}',0,$${text}$$);" + )) + .execute(&pool) + .await?; + let (model, server) = match replies { + Some(replies) => { + let model = Model { + replies: Arc::new(Mutex::new(replies)), + requests: Arc::new(Mutex::new(Vec::new())), + }; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let router = Router::new() + .route("/chat/completions", post(reply)) + .with_state(model.clone()); + let server = + tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + utopia_store::settings::upsert( + &pool, + ws, + Some(&endpoint), + None, + Some("scripted"), + None, + None, + None, + None, + ) + .await?; + (Some(model), Some(server)) + } + None => (None, None), + }; + let dir = tempfile::tempdir()?; + let cfg = utopia_core::config::AppConfig { + data_dir: dir.path().to_string_lossy().into_owned(), + ..Default::default() + }; + let search = Arc::new(utopia_search::SearchIndex::open( + &dir.path().join("search"), + )?); + let state = AppState::new(pool.clone(), &cfg, search, "test-only".into()); + // 抽取：推送路径按契约解析、不问模型（0054）；这一步把陈述与时间词写进库 + crate::extraction::extract_document(&state, doc, utopia_core::models::Proposer::default()) + .await?; + Ok(Some(Self { + pool, + state, + org, + kb, + doc, + model, + server, + _dir: dir, + })) + } + + /// 陈述的开放行，按宾语名排 + async fn fact_rows(&self) -> anyhow::Result> { + Ok(sqlx::query_as( + "SELECT o.canonical_name, f.valid_from, f.valid_from_precision, + f.valid_to, f.valid_to_precision, f.valid_from_grade + FROM facts f JOIN entities o ON o.id = f.object_id + WHERE f.kb_id = $1 AND f.layer = 'open' AND f.invalidated_at IS NULL + ORDER BY o.canonical_name", + ) + .bind(self.kb) + .fetch_all(&self.pool) + .await?) + } + + /// 时间提及行：(字、槽、等级、解释形状) + async fn mention_rows(&self) -> anyhow::Result> { + Ok(sqlx::query_as( + "SELECT text, role, grade, shape FROM time_mentions + WHERE kb_id = $1 ORDER BY role, text", + ) + .bind(self.kb) + .fetch_all(&self.pool) + .await?) + } + + async fn requests(&self) -> Vec { + match &self.model { + Some(m) => m.requests.lock().unwrap().clone(), + None => Vec::new(), + } + } + + async fn cleanup(self) -> anyhow::Result<()> { + if let Some(server) = self.server { + server.abort(); + } + sqlx::query("DELETE FROM jobs WHERE payload->>'document_id'=$1 OR payload->>'kb_id'=$2") + .bind(self.doc.to_string()) + .bind(self.kb.to_string()) + .execute(&self.pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id=$1") + .bind(self.org) + .execute(&self.pool) + .await?; + Ok(()) + } +} + +fn payload(when: &str, ended: Option<&str>) -> Value { + json!({ + "external_id": "obs-1", + "doc_time": "2026-10-01T00:00:00Z", + "e": [["cup-7", "cup", true], ["kitchen table", "table", true]], + "s": [[null, "cup-7", "is on", "kitchen table", null, null, when, ended]], + "n": [] + }) +} + +fn payload_two(when_a: &str, when_b: &str) -> Value { + json!({ + "external_id": "obs-1", + "doc_time": "2026-10-01T00:00:00Z", + "e": [["cup-7", "cup", true], ["kitchen table", "table", true], ["counter", "table", true]], + "s": [ + [null, "cup-7", "is on", "kitchen table", null, null, when_a, null], + [null, "cup-7", "is on", "counter", null, null, when_b, null] + ], + "n": [] + }) +} + +fn day(y: i32, m: u32, d: u32) -> Option> { + Some(Utc.with_ymd_and_hms(y, m, d, 0, 0, 0).unwrap()) +} + +fn at(y: i32, mo: u32, d: u32, h: u32, mi: u32, s: u32) -> Option> { + Some(Utc.with_ymd_and_hms(y, mo, d, h, mi, s).unwrap()) +} + +/// 没有对话模型：载荷写明的日期由代码读成 A 级、字写到的精度，任务照常完成； +/// 字本身（`"2026-09-23"`）留着当提及与证据 +#[tokio::test] +async fn a_pushed_date_is_resolved_in_code_without_a_model() -> anyhow::Result<()> { + let Some(rig) = Rig::new(payload("2026-09-23", None), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + day(2026, 9, 23), + Some("day".to_string()), + day(2026, 9, 23), + Some("day".to_string()), + Some("A".to_string()), + )], + "a bare date is a single date at day precision, grade A" + ); + assert_eq!( + rig.mention_rows().await?, + vec![( + "2026-09-23".to_string(), + "when".to_string(), + Some("A".to_string()), + Some("point".to_string()) + )], + "the literal string stays as the mention" + ); + rig.cleanup().await +} + +/// 带 `Z` 与带偏移的时刻：精度到秒，值是换算成 UTC 的同一刻——偏移不是贴标签， +/// 是真的进了值 +#[tokio::test] +async fn a_pushed_timestamp_resolves_at_second_precision_and_its_zone_counts() -> anyhow::Result<()> +{ + let Some(rig) = Rig::new(payload("2026-09-23T08:14:03Z", None), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + at(2026, 9, 23, 8, 14, 3), + Some("second".to_string()), + at(2026, 9, 23, 8, 14, 3), + Some("second".to_string()), + Some("A".to_string()), + )] + ); + rig.cleanup().await?; + + let Some(rig) = Rig::new(payload("2026-09-23T08:14:03+08:00", None), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + // +08:00 的 08:14:03 是 UTC 的 00:14:03 + at(2026, 9, 23, 0, 14, 3), + Some("second".to_string()), + at(2026, 9, 23, 0, 14, 3), + Some("second".to_string()), + Some("A".to_string()), + )] + ); + rig.cleanup().await +} + +/// 字写到哪一级，精度就到哪一级：`T08` 是小时、`T08:14` 是分钟——不替它补 +/// 从没写过的秒（0024） +#[tokio::test] +async fn the_precision_is_the_one_the_payload_writes() -> anyhow::Result<()> { + for (when, expected, precision) in [ + ("2026-09-23", (0u32, 0u32, 0u32), "day"), + ("2026-09-23T08Z", (8, 0, 0), "hour"), + ("2026-09-23T08:14Z", (8, 14, 0), "minute"), + ("2026-09-23T08:14:03Z", (8, 14, 3), "second"), + ] { + let Some(rig) = Rig::new(payload(when, None), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + let (h, m, s) = expected; + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + at(2026, 9, 23, h, m, s), + Some(precision.to_string()), + at(2026, 9, 23, h, m, s), + Some(precision.to_string()), + Some("A".to_string()), + )], + "{when} resolves at {precision}" + ); + rig.cleanup().await?; + } + Ok(()) +} + +/// `when` 加 `ended` 是一个区间：起是 when、止是 ended，各带自己的精度 +#[tokio::test] +async fn when_plus_ended_is_an_interval() -> anyhow::Result<()> { + let Some(rig) = Rig::new(payload("2026-09-20", Some("2026-09-23T08:14:03Z")), None).await? + else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + day(2026, 9, 20), + Some("day".to_string()), + at(2026, 9, 23, 8, 14, 3), + Some("second".to_string()), + Some("A".to_string()), + )] + ); + rig.cleanup().await +} + +/// 光杆的年与年月不在这一刀里：没有模型时留着 C，任务照常收工； +/// 自然语言的时间词同样留着 C +#[tokio::test] +async fn a_bare_year_or_year_month_or_words_are_left_for_a_model() -> anyhow::Result<()> { + for when in ["2026", "2026-09", "last Tuesday"] { + let Some(rig) = Rig::new(payload(when, None), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.mention_rows().await?, + vec![( + when.to_string(), + "when".to_string(), + Some("C".to_string()), + None + )], + "{when:?} is not code-resolved" + ); + assert_eq!( + rig.fact_rows().await?, + vec![( + "kitchen table".to_string(), + None, + None, + None, + None, + Some("C".to_string()), + )], + "{when:?}: the statement keeps its time words, unresolved, and the job ends" + ); + rig.cleanup().await?; + } + Ok(()) +} + +/// 一份载荷两种字：读得动的解成 A，读不动的留 C——一篇文档里两种结局可以并存， +/// 任务照常收工 +#[tokio::test] +async fn a_mixed_push_resolves_what_code_can_and_leaves_the_rest() -> anyhow::Result<()> { + let Some(rig) = Rig::new(payload_two("2026-09-23", "last Tuesday"), None).await? else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + assert_eq!( + rig.mention_rows().await?, + vec![ + ( + "2026-09-23".to_string(), + "when".to_string(), + Some("A".to_string()), + Some("point".to_string()) + ), + ( + "last Tuesday".to_string(), + "when".to_string(), + Some("C".to_string()), + None + ), + ] + ); + rig.cleanup().await +} + +/// 文档里的字长得像 RFC 3339 也不走推送那条路：判据是来源种类，不是字的形状。 +/// 配了模型时这条提及照样送出去问——来源不是 `statements`，代码不代读 +#[tokio::test] +async fn a_document_mention_that_looks_machine_readable_stays_on_the_model_path( +) -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = sqlx::PgPool::connect(&url).await?; + utopia_store::db::migrate(&pool).await?; + let (org, ws, kb, doc, chunk) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + let text = "The cup was on the kitchen table at 2026-09-23T08:14:03Z."; + sqlx::raw_sql(&format!( + "INSERT INTO organizations(id,name) VALUES ('{org}','document-path'); + INSERT INTO workspaces(id,org_id,name) VALUES ('{ws}','{org}','document-path'); + INSERT INTO knowledge_bases(id,workspace_id,name) VALUES ('{kb}','{ws}','document-path'); + INSERT INTO documents(id,kb_id,filename,sha256) VALUES ('{doc}','{kb}','note.md','x{doc}'); + INSERT INTO chunks(id,kb_id,document_id,seq,text) VALUES + ('{chunk}','{kb}','{doc}',0,$${text}$$);" + )) + .execute(&pool) + .await?; + let absolute = json!({ + "kind": "absolute", + // 契约里日期部件用紧凑键（y/m/d/h/min/s） + "from": {"y": 2026, "m": 9, "d": 23, "h": 8, "min": 14, "s": 3} + }); + let model = Model { + replies: Arc::new(Mutex::new(vec![ + json!({ + "e": [["the cup", "cup", 1], ["the kitchen table", "table", 1]], + "s": [["The cup was on the kitchen table at 2026-09-23T08:14:03Z.", + "the cup", "was on", "the kitchen table", null, null, + "2026-09-23T08:14:03Z", null]], + "n": [], + "t": [] + }) + .to_string(), + // 解释：模型照答——走的还是模型这条路，代码没代读 + json!({ "m": [[0, "point", absolute, "second"]] }).to_string(), + ])), + requests: Arc::new(Mutex::new(Vec::new())), + }; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let router = Router::new() + .route("/chat/completions", post(reply)) + .with_state(model.clone()); + let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + utopia_store::settings::upsert( + &pool, + ws, + Some(&endpoint), + None, + Some("scripted"), + None, + None, + None, + None, + ) + .await?; + let dir = tempfile::tempdir()?; + let cfg = utopia_core::config::AppConfig { + data_dir: dir.path().to_string_lossy().into_owned(), + ..Default::default() + }; + let search = Arc::new(utopia_search::SearchIndex::open( + &dir.path().join("search"), + )?); + let state = AppState::new(pool.clone(), &cfg, search, "test-only".into()); + crate::extraction::extract_document(&state, doc, utopia_core::models::Proposer::default()) + .await?; + resolve_document(&state, doc).await?; + // 一次抽取 + 一次解释：长得像日期的字照样问过模型；代码没有替它读 + assert_eq!( + model.requests.lock().unwrap().len(), + 2, + "the mention went to the model: source identity, not string shape" + ); + // 陈述行：宾语是「the kitchen table」的那条（名字事实等别的行不算） + let row: (Option>, Option) = sqlx::query_as( + "SELECT f.valid_from, f.valid_from_grade FROM facts f + JOIN entities o ON o.id = f.object_id + WHERE f.kb_id = $1 AND f.layer = 'open' AND o.canonical_name = 'the kitchen table'", + ) + .bind(kb) + .fetch_one(&pool) + .await?; + assert_eq!(row, (at(2026, 9, 23, 8, 14, 3), Some("A".to_string()))); + server.abort(); + sqlx::query("DELETE FROM jobs WHERE payload->>'document_id'=$1 OR payload->>'kb_id'=$2") + .bind(doc.to_string()) + .bind(kb.to_string()) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id=$1") + .bind(org) + .execute(&pool) + .await?; + Ok(()) +} + +/// 配了模型也一样：代码读得动的提及不送模型（回复里那条 `anchored` 是模型对 +/// 「上周」的答法——它锚在另一条提及上，而那一条正是代码读的）。再跑一次 +/// `resolve_time` 答案不变：时间戳照样不问，「上周」重问一回得的也是同一个位置—— +/// 代码读的没有模型可盖 +#[tokio::test] +async fn a_code_resolved_mention_is_never_asked_and_survives_reruns() -> anyhow::Result<()> { + // 两个时间词：代码能读的时间戳是第 0 条distinct；「上周」是第 1 条，只有它去问模型。 + // 模型把它锚在第 0 条上——代码读出的解释与模型给的解释在同一张锚点表里 + let last_week = json!({ "m": [[1, "point", { + "kind": "anchored", "anchor": {"kind": "mention", "id": 0}, + "offset": {"count": 1, "unit": "week", "direction": "before"} + }, "week"]] }) + .to_string(); + let Some(rig) = Rig::new( + payload_two("2026-09-23T08:14:03Z", "上周"), + // 重跑还会再替「上周」问一次：给它同一个答案 + Some(vec![last_week.clone(), last_week]), + ) + .await? + else { + return Ok(()); + }; + resolve_document(&rig.state, rig.doc).await?; + // 抽取不问模型（推送路），这里只有一次解释调用：只问「上周」。 + // 提示词的「提及」一栏里没有那条时间戳——它会作为载荷正文（sentence）出现在 + // 「上周」的语境里，那是文档自己的字；关键是它没有被当成一条提及去问 + let requests = rig.requests().await; + assert_eq!(requests.len(), 1, "only the interpretation call"); + let asked = requests[0].to_string(); + // 「提及」一栏只有「上周」这一条被问（- id 只出现一次），时间戳不在其中—— + // 它会作为载荷正文出现在「上周」的语境里，那是文档自己的字 + assert!(asked.contains("上周"), "asked: {asked}"); + assert_eq!(asked.matches("- id ").count(), 1, "asked: {asked}"); + + let first = rig.fact_rows().await?; + assert_eq!( + first, + vec![ + // 「上周」锚在 9-23T08:14:03Z 上前一周：2026-09-16，日精度，B 级 + ( + "counter".to_string(), + day(2026, 9, 16), + Some("day".to_string()), + day(2026, 9, 16), + Some("day".to_string()), + Some("B".to_string()), + ), + ( + "kitchen table".to_string(), + at(2026, 9, 23, 8, 14, 3), + Some("second".to_string()), + at(2026, 9, 23, 8, 14, 3), + Some("second".to_string()), + Some("A".to_string()), + ), + ], + "code-resolved, precision as written" + ); + // 重跑：同样的答案；脚本端点有第二份回复等着「上周」 + resolve_document(&rig.state, rig.doc).await?; + let requests = rig.requests().await; + assert_eq!( + requests.len(), + 2, + "the rerun asked the model for 上周 again" + ); + assert_eq!( + requests[1].to_string().matches("- id ").count(), + 1, + "the rerun also asked only about 上周" + ); + assert_eq!( + rig.fact_rows().await?, + first, + "a rerun gives the same answer" + ); + rig.cleanup().await +} diff --git a/crates/utopia-store/src/governance.rs b/crates/utopia-store/src/governance.rs index 0eb8af84a..2317c264f 100644 --- a/crates/utopia-store/src/governance.rs +++ b/crates/utopia-store/src/governance.rs @@ -843,7 +843,7 @@ pub async fn namesakes( let rows: Vec = sqlx::query_as( "SELECT e.canonical_name AS name, t.label AS type_label, (SELECT count(*) FROM facts f - WHERE f.invalidated_at IS NULL + WHERE f.kb_id = e.kb_id AND f.invalidated_at IS NULL AND (f.subject_id = e.id OR f.object_id = e.id)) AS facts, (e.merged_into IS NOT NULL) AS merged FROM entities e diff --git a/crates/utopia-store/src/graph.rs b/crates/utopia-store/src/graph.rs index ec77e8557..d2bf3d82a 100644 --- a/crates/utopia-store/src/graph.rs +++ b/crates/utopia-store/src/graph.rs @@ -1274,7 +1274,7 @@ fn node_sql(as_of: Option, owner: Option) -> String { coalesce(t.shape, 'circle') AS shape, e.disambiguator, (SELECT count(*) FROM facts f - WHERE ({subject} = e.id OR {object} = e.id) AND {held} AND {not_name}) AS degree + WHERE f.kb_id = e.kb_id AND ({subject} = e.id OR {object} = e.id) AND {held} AND {not_name}) AS degree FROM entities e LEFT JOIN entity_types t ON t.id = e.type_id" ) } diff --git a/crates/utopia-store/src/ontology.rs b/crates/utopia-store/src/ontology.rs index 225cabcf8..87f4c11a4 100644 --- a/crates/utopia-store/src/ontology.rs +++ b/crates/utopia-store/src/ontology.rs @@ -19,9 +19,15 @@ pub async fn entity_type_views(pool: &PgPool, kb_id: Uuid) -> AppResult = sqlx::query_as( "SELECT e.id, e.canonical_name AS name, (SELECT count(*) FROM facts f - WHERE (f.subject_id = e.id OR f.object_id = e.id) - AND f.invalidated_at IS NULL) AS fact_count + WHERE f.kb_id = e.kb_id + AND (f.subject_id = e.id OR f.object_id = e.id) + AND f.invalidated_at IS NULL) AS fact_count FROM entities e WHERE e.kb_id = $1 AND e.type_id = $2 AND e.merged_into IS NULL ORDER BY lower(e.canonical_name), e.id @@ -85,9 +92,15 @@ pub async fn relation_type_views(pool: &PgPool, kb_id: Uuid) -> AppResult, +) -> AppResult { + Ok(sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM relation_types + WHERE kb_id = $1 AND embedding IS NOT NULL + AND ($2::text IS NULL OR kind = $2))", + ) + .bind(kb_id) + .bind(only_kind) + .fetch_one(pool) + .await?) +} + /// 一次插完一批类，返回 key → id。 /// /// **存在的理由是 fsync。** 逐条 `execute(pool)` 每条各自提交，导入 schema.org @@ -2496,6 +2531,27 @@ pub struct AgentProposalReport { pub rejected: i64, } +/// 文档说了、本体还放不下的：没绑到类的类别词，和没绑到属性的短语形状（判成 none 的 +/// 与两票不一致的）。本体代理读的就是这两样（0061 决定 2）；工作台拿这两个数说「还缺什么」， +/// 不再拿 0003 的 `ontology_misses`——开放图谱的抽取不往那张表里写 +#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow, serde::Serialize)] +pub struct Uncovered { + pub kind_words: i64, + pub phrases: i64, +} + +pub async fn uncovered(pool: &PgPool, kb_id: Uuid) -> AppResult { + Ok(sqlx::query_as( + "SELECT (SELECT count(*) FROM type_bindings + WHERE kb_id = $1 AND status <> 'bound') AS kind_words, + (SELECT count(*) FROM phrase_bindings + WHERE kb_id = $1 AND status <> 'bound') AS phrases", + ) + .bind(kb_id) + .fetch_one(pool) + .await?) +} + pub async fn agent_proposal_report(pool: &PgPool, kb_id: Uuid) -> AppResult { Ok(sqlx::query_as( "SELECT count(*) FILTER (WHERE status = 'open') AS open, diff --git a/crates/utopia-store/src/phrase_bindings.rs b/crates/utopia-store/src/phrase_bindings.rs index f658aea02..b4c1fe9c2 100644 --- a/crates/utopia-store/src/phrase_bindings.rs +++ b/crates/utopia-store/src/phrase_bindings.rs @@ -73,37 +73,34 @@ pub async fn signatures(pool: &PgPool, kb_id: Uuid) -> AppResult '' ), - grouped AS ( - SELECT phrase, subject_type_id, object_type_id, object_is_value, count(*) AS count - FROM live GROUP BY phrase, subject_type_id, object_type_id, object_is_value - ), - picked AS ( - SELECT g.phrase, g.subject_type_id, g.object_type_id, g.object_is_value, - l.subject_name || ' —' || l.spelling || '→ ' || l.object_name AS example, - coalesce(substr(c.text, l.quote_start + 1, l.quote_end - l.quote_start), '') AS quote, - row_number() OVER (PARTITION BY g.phrase, g.subject_type_id, g.object_type_id, g.object_is_value + -- 每组的计数和前 3 条例句在一趟里聚出来。从前每组的例句、引文各是一个关联子查询， + -- 每个都把全部陈述重扫一遍，耗时随开放陈述数平方增长：3.6 万条时一趟约 200 秒（#1096）。 + -- 计数也在这里取，不另聚一份再按四个键用 IS NOT DISTINCT FROM 接回来——那种连接 + -- Postgres 做不了哈希；分块的原文只为留下的那三条去读 + ranked AS ( + SELECT l.*, + row_number() OVER (PARTITION BY l.phrase, l.subject_type_id, l.object_type_id, l.object_is_value ORDER BY l.recorded_at, l.id) AS rn - FROM grouped g - JOIN live l ON l.phrase = g.phrase - AND l.subject_type_id IS NOT DISTINCT FROM g.subject_type_id - AND l.object_type_id IS NOT DISTINCT FROM g.object_type_id - AND l.object_is_value = g.object_is_value - LEFT JOIN chunks c ON c.id = l.chunk_id + FROM live l + ), + grouped AS ( + SELECT r.phrase, r.subject_type_id, r.object_type_id, r.object_is_value, count(*) AS count, + array_agg(r.subject_name || ' —' || r.spelling || '→ ' || r.object_name ORDER BY r.rn) + FILTER (WHERE r.rn <= 3) AS examples, + array_agg(coalesce(substr(c.text, r.quote_start + 1, r.quote_end - r.quote_start), '') ORDER BY r.rn) + FILTER (WHERE r.rn <= 3) AS quotes + FROM ranked r + LEFT JOIN chunks c ON c.id = r.chunk_id AND r.rn <= 3 + GROUP BY r.phrase, r.subject_type_id, r.object_type_id, r.object_is_value ) SELECT g.phrase, g.subject_type_id, st.key AS subject_type_key, g.object_type_id, ot.key AS object_type_key, g.object_is_value, g.count, - ARRAY(SELECT p.example FROM picked p - WHERE p.phrase = g.phrase AND p.subject_type_id IS NOT DISTINCT FROM g.subject_type_id - AND p.object_type_id IS NOT DISTINCT FROM g.object_type_id - AND p.object_is_value = g.object_is_value AND p.rn <= 3 ORDER BY p.rn) AS examples, - ARRAY(SELECT p.quote FROM picked p - WHERE p.phrase = g.phrase AND p.subject_type_id IS NOT DISTINCT FROM g.subject_type_id - AND p.object_type_id IS NOT DISTINCT FROM g.object_type_id - AND p.object_is_value = g.object_is_value AND p.rn <= 3 ORDER BY p.rn) AS quotes + coalesce(g.examples, '{{}}') AS examples, coalesce(g.quotes, '{{}}') AS quotes FROM grouped g LEFT JOIN entity_types st ON st.id = g.subject_type_id LEFT JOIN entity_types ot ON ot.id = g.object_type_id - ORDER BY g.count DESC, g.phrase", + -- 计数和短语相同、两端的类不同的签名，顺序也得定下来：对齐按这个顺序分批送模型 + ORDER BY g.count DESC, g.phrase, g.subject_type_id, g.object_type_id, g.object_is_value", phrase = phrase_sql("f.phrase") ); Ok(sqlx::query_as(&sql).bind(kb_id).fetch_all(pool).await?) diff --git a/crates/utopia-store/src/resolution.rs b/crates/utopia-store/src/resolution.rs index d17854a0b..54f564268 100644 --- a/crates/utopia-store/src/resolution.rs +++ b/crates/utopia-store/src/resolution.rs @@ -256,9 +256,49 @@ pub async fn resolve_mention( exclude: &[Uuid], ) -> AppResult { let mut r = resolve_by_name(pool, kb_id, type_id, raw_name, context, text, exclude).await?; - let Some(query) = name_vector else { - return Ok(r); - }; + if let Some(query) = name_vector { + propose_by_name_vector(pool, kb_id, type_id, raw_name, query, exclude, &mut r).await?; + } + // 判过「不是一个」的对不再排。新建的实体没有过去，没东西要排的也不用问 + if !r.created && !r.reviews.is_empty() { + let kept = kept_apart_from(pool, kb_id, r.entity_id).await?; + r.reviews.retain(|v| !kept.contains(&v.other_id)); + } + Ok(r) +} + +/// 跟这个实体判过「不是一个」的那些实体。 +/// +/// **每次提到都会走到召回**：名字向量通道不分新建还是归并，同名易混类型那条路也是 +/// 归并之后照常入队。而审核表的唯一索引只管 pending 的行——判完的一对下次被提到， +/// 就又是一条新的 pending。一个有六万实体的库上，`migration 395` 每出现在一篇新文档里， +/// 它跟 `migration 394` 等八个近邻就重新排一遍队，裁决器（或人）分开过多少次都一样 +/// （#1104）。分开这个决定要记得住，[`crate::names::pair_shared_name`] 早就是这么做的。 +/// +/// 人仍然可以在审核页把分开过的一对合并：这里拦的是自动重提，不是合并本身。 +async fn kept_apart_from(pool: &PgPool, kb_id: Uuid, entity_id: Uuid) -> AppResult> { + let rows: Vec<(Uuid,)> = sqlx::query_as( + "SELECT CASE WHEN left_id = $2 THEN right_id ELSE left_id END + FROM resolution_reviews + WHERE kb_id = $1 AND status = 'kept' AND (left_id = $2 OR right_id = $2)", + ) + .bind(kb_id) + .bind(entity_id) + .fetch_all(pool) + .await?; + Ok(rows.into_iter().map(|(id,)| id).collect()) +} + +/// 通道 2：mention 名字向量的近邻，各排一对（规矩见 [`resolve_mention`]） +async fn propose_by_name_vector( + pool: &PgPool, + kb_id: Uuid, + type_id: Option, + raw_name: &str, + query: &[f32], + exclude: &[Uuid], + r: &mut Resolution, +) -> AppResult<()> { let mention_name = normalize_name(raw_name).to_lowercase(); let mention_family = match type_id { Some(t) => type_label(pool, t) @@ -298,7 +338,7 @@ pub async fn resolve_mention( stage: ReviewStage::Adjudicating, }); } - Ok(r) + Ok(()) } async fn type_label(pool: &PgPool, type_id: Uuid) -> AppResult> { @@ -337,7 +377,7 @@ async fn resolve_by_name( let candidates: Vec = sqlx::query_as(&format!( "SELECT e.id, e.canonical_name, e.profile_embedding, e.profile_n, (SELECT count(*) FROM facts f - WHERE (f.subject_id = e.id OR f.object_id = e.id) + WHERE f.kb_id = e.kb_id AND (f.subject_id = e.id OR f.object_id = e.id) AND f.invalidated_at IS NULL AND {not_name}) AS degree FROM entities e -- IS NOT DISTINCT FROM 而不是 =（0009 的那个陷阱）：开放图谱里的实体都没有类 @@ -1254,7 +1294,7 @@ pub async fn existing_by_name( AND e.description IS NULL AND (lower(e.canonical_name) = ANY($2) OR {named}) ORDER BY (SELECT count(*) FROM facts f - WHERE (f.subject_id = e.id OR f.object_id = e.id) AND {not_name}) DESC, + WHERE f.kb_id = e.kb_id AND (f.subject_id = e.id OR f.object_id = e.id) AND f.invalidated_at IS NULL AND {not_name}) DESC, e.created_at LIMIT 1", named = crate::names::has_name_in("e", 1, 2), @@ -1398,7 +1438,7 @@ async fn review_side(pool: &PgPool, kb_id: Uuid, entity_id: Uuid) -> AppResult AppResult AppResult<(Uuid, Uuid)> { + // 事实索引以 kb_id 开头；先用标量子查询取库，让 OR 两侧都能按复合索引定位。 + // PG16 上直接 JOIN 可能只按 kb_id 扫描整个库，再过滤 subject/object。 let (deg_a,): (i64,) = sqlx::query_as( - "SELECT count(*) FROM facts WHERE (subject_id = $1 OR object_id = $1) AND invalidated_at IS NULL", + "SELECT count(*) FROM facts + WHERE kb_id = (SELECT kb_id FROM entities WHERE id = $1) + AND (subject_id = $1 OR object_id = $1) AND invalidated_at IS NULL", ) .bind(a) .fetch_one(pool) .await?; let (deg_b,): (i64,) = sqlx::query_as( - "SELECT count(*) FROM facts WHERE (subject_id = $1 OR object_id = $1) AND invalidated_at IS NULL", + "SELECT count(*) FROM facts + WHERE kb_id = (SELECT kb_id FROM entities WHERE id = $1) + AND (subject_id = $1 OR object_id = $1) AND invalidated_at IS NULL", ) .bind(b) .fetch_one(pool) diff --git a/crates/utopia-store/src/type_bindings.rs b/crates/utopia-store/src/type_bindings.rs index 5c6859772..e83fe701e 100644 --- a/crates/utopia-store/src/type_bindings.rs +++ b/crates/utopia-store/src/type_bindings.rs @@ -67,28 +67,42 @@ pub async fn signatures<'e>( WHERE e.kb_id = $1 AND e.merged_into IS NULL AND e.specific_type IS NOT NULL AND btrim(e.specific_type) <> '' ), + -- 写法、例名、关系短语的前 3 个各在一趟里聚出来。从前每个词各跑三个关联子查询， + -- 每个都把全部实体（短语那个还连着陈述）重扫一遍，耗时随库的大小平方增长（#1096） grouped AS ( - SELECT kind_word, count(*) AS count FROM live GROUP BY kind_word + SELECT kind_word, count(*) AS count, + (array_agg(canonical_name ORDER BY (description IS NOT NULL), created_at, id))[1:3] AS examples + FROM live GROUP BY kind_word + ), + spellings AS ( + SELECT kind_word, spelling, + row_number() OVER (PARTITION BY kind_word ORDER BY count(*) DESC, spelling) AS rn + FROM live GROUP BY kind_word, spelling + ), + words AS ( + SELECT kind_word, array_agg(spelling ORDER BY rn) AS words + FROM spellings WHERE rn <= 3 GROUP BY kind_word + ), + phrase_counts AS ( + SELECT l.kind_word, f.phrase, + row_number() OVER (PARTITION BY l.kind_word ORDER BY count(*) DESC, f.phrase) AS rn + FROM facts f JOIN live l ON l.id = f.subject_id + WHERE f.kb_id = $1 AND f.layer = 'open' AND f.invalidated_at IS NULL + AND f.phrase IS NOT NULL + GROUP BY l.kind_word, f.phrase + ), + phrases AS ( + SELECT kind_word, array_agg(phrase ORDER BY rn) AS phrases + FROM phrase_counts WHERE rn <= 3 GROUP BY kind_word ) SELECT g.kind_word, g.count, - ARRAY(SELECT s.spelling FROM ( - SELECT l.spelling, count(*) AS n FROM live l - WHERE l.kind_word = g.kind_word - GROUP BY l.spelling ORDER BY n DESC, l.spelling LIMIT 3) s - ) AS words, - ARRAY(SELECT l.canonical_name FROM live l - WHERE l.kind_word = g.kind_word - ORDER BY (l.description IS NOT NULL), l.created_at, l.id LIMIT 3 - ) AS examples, - ARRAY(SELECT s.phrase FROM ( - SELECT f.phrase, count(*) AS n - FROM facts f JOIN live l ON l.id = f.subject_id - WHERE f.kb_id = $1 AND f.layer = 'open' AND f.invalidated_at IS NULL - AND f.phrase IS NOT NULL AND l.kind_word = g.kind_word - GROUP BY f.phrase ORDER BY n DESC, f.phrase LIMIT 3) s - ) AS phrases + coalesce(w.words, '{{}}') AS words, + g.examples, + coalesce(p.phrases, '{{}}') AS phrases FROM grouped g + LEFT JOIN words w ON w.kind_word = g.kind_word + LEFT JOIN phrases p ON p.kind_word = g.kind_word ORDER BY g.count DESC, g.kind_word", kind = kind_word_sql("e.specific_type") ); diff --git a/crates/utopia-store/tests/store/a_similar_name_is_proposed_not_merged.rs b/crates/utopia-store/tests/store/a_similar_name_is_proposed_not_merged.rs index 2af38290b..42f32f551 100644 --- a/crates/utopia-store/tests/store/a_similar_name_is_proposed_not_merged.rs +++ b/crates/utopia-store/tests/store/a_similar_name_is_proposed_not_merged.rs @@ -332,3 +332,77 @@ async fn a_new_name_is_pending_until_its_vector_is_set() -> anyhow::Result<()> { teardown(&pool, &f).await?; run } + +/// 判过「不是一个」的一对，实体再被提到时不再排（#1104）。名字向量通道每次提到都跑， +/// 而审核表只对 pending 的行去重：从前裁决器分开过的一对，下一篇文档一提到就又是一条 +/// 新的 pending。还没判的那一对不受影响，照常提议 +#[tokio::test] +async fn a_pair_kept_apart_is_not_proposed_again() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + let f = seed(&pool, "name-vector-kept-apart").await?; + let run = async { + let near: Vec = vec![0.95, 0.31, 0.0]; + let mention = || { + utopia_store::resolution::resolve_mention( + &pool, + f.kb, + Some(f.device), + "海探1", + None, + Some(&near), + None, + &[], + ) + }; + let first = mention().await?; + assert!(first.created); + for v in &first.reviews { + utopia_store::resolution::create_review( + &pool, + f.kb, + first.entity_id, + v.other_id, + v.score, + &v.reason, + v.stage, + ) + .await?; + } + // 裁决器判 海探1 与 海洋探测器1号 不是一个；跟 海洋探测队长 的那一对还没判 + let (review,): (Uuid,) = sqlx::query_as( + "SELECT id FROM resolution_reviews WHERE kb_id = $1 AND status = 'pending' + AND $2 IN (left_id, right_id)", + ) + .bind(f.kb) + .bind(f.probe) + .fetch_one(&pool) + .await?; + utopia_store::resolution::close_review_auto(&pool, review, "kept", "kept_apart|test") + .await?; + + let again = mention().await?; + assert!(!again.created, "同一个字面名字归到上一轮建的实体"); + assert_eq!(again.entity_id, first.entity_id); + let proposed: Vec = vector_reviews(&again) + .into_iter() + .map(|(id, _)| id) + .collect(); + assert!( + !proposed.contains(&f.probe), + "分开过的一对不该再提议：{:?}", + again.reviews + ); + assert!( + proposed.contains(&f.captain), + "还没判的一对照常提议（入队时对 pending 幂等）：{:?}", + again.reviews + ); + Ok::<_, anyhow::Error>(()) + } + .await; + teardown(&pool, &f).await?; + run +} diff --git a/crates/utopia-store/tests/store/an_empty_ontology_does_not_cover_the_corpus.rs b/crates/utopia-store/tests/store/an_empty_ontology_does_not_cover_the_corpus.rs new file mode 100644 index 000000000..b8787100b --- /dev/null +++ b/crates/utopia-store/tests/store/an_empty_ontology_does_not_cover_the_corpus.rs @@ -0,0 +1,98 @@ +//! 文档说了、本体还放不下的有多少：没绑到类的类别词、没绑到属性的短语形状。 +//! +//! 工作台从前拿 0003 的 `ontology_misses` 说「本体覆盖了语料」，而开放图谱的抽取不往那张表 +//! 里写：一个一个类都没有的库，二十二个类别词没着落，页面上写的是「覆盖了」。这两个数是 +//! 本体代理读的那两样（0061 决定 2），别的库的不算。 +//! +//! 没有 `UTOPIA_DATABASE_URL` 时跳过而不是失败。自建自拆，绝不碰已有的库。 + +use sqlx::PgPool; +use utopia_store::ontology::{uncovered, Uncovered}; +use uuid::Uuid; + +const ORG: &str = "empty-ontology-does-not-cover-test"; + +#[tokio::test] +async fn kind_words_and_phrases_with_no_place_are_counted_per_base() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + sqlx::query("DELETE FROM organizations WHERE name = $1") + .bind(ORG) + .execute(&pool) + .await?; + let (org, ws, kb, other, person, leads) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + sqlx::raw_sql(&format!( + "INSERT INTO organizations(id,name) VALUES ('{org}','{ORG}'); + INSERT INTO workspaces(id,org_id,name) VALUES ('{ws}','{org}','{ORG}'); + INSERT INTO knowledge_bases(id,workspace_id,name) VALUES ('{kb}','{ws}','a'), ('{other}','{ws}','b');" + )) + .execute(&pool) + .await?; + + let run = async { + // 什么都还没读过：没有要等的 + let none = Uncovered { + kind_words: 0, + phrases: 0, + }; + assert_eq!(uncovered(&pool, kb).await?, none); + + // 读过文档、一个类都没有：每个类别词都没着落 + for word in ["person", "project", "city"] { + sqlx::query( + "INSERT INTO type_bindings (id, kb_id, kind_word, status) VALUES ($1, $2, $3, 'none')", + ) + .bind(Uuid::now_v7()) + .bind(kb) + .bind(word) + .execute(&pool) + .await?; + } + assert_eq!(uncovered(&pool, kb).await?.kind_words, 3); + + // 加了一个类、一条属性：绑上的不算，判成 none 的和两票不一致的都算 + sqlx::raw_sql(&format!( + "INSERT INTO entity_types(id,kb_id,key,label) VALUES ('{person}','{kb}','person','person'); + INSERT INTO relation_types(id,kb_id,key,label,kind,temporal,description) + VALUES ('{leads}','{kb}','leads','leads','relation','state',''); + UPDATE type_bindings SET status='bound', type_id='{person}' + WHERE kb_id='{kb}' AND kind_word='person'; + UPDATE type_bindings SET status='undecided' WHERE kb_id='{kb}' AND kind_word='city'; + INSERT INTO phrase_bindings (id, kb_id, phrase, subject_type_id, status, relation_type_id, direction) + VALUES ('{}','{kb}','负责','{person}','bound','{leads}','forward'), + ('{}','{kb}','汇报给','{person}','none',NULL,NULL), + ('{}','{kb}','调任','{person}','undecided',NULL,NULL); + INSERT INTO type_bindings (id, kb_id, kind_word, status) VALUES ('{}','{other}','supplier','none');", + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + )) + .execute(&pool) + .await?; + assert_eq!( + uncovered(&pool, kb).await?, + Uncovered { + kind_words: 2, + phrases: 2, + } + ); + assert_eq!(uncovered(&pool, other).await?.kind_words, 1, "each base counts its own"); + anyhow::Ok(()) + } + .await; + sqlx::query("DELETE FROM organizations WHERE id = $1") + .bind(org) + .execute(&pool) + .await?; + run +} diff --git a/crates/utopia-store/tests/store/main.rs b/crates/utopia-store/tests/store/main.rs index 565ef2f3b..029c0ba20 100644 --- a/crates/utopia-store/tests/store/main.rs +++ b/crates/utopia-store/tests/store/main.rs @@ -86,6 +86,7 @@ mod a_search_puts_the_exact_name_first; mod a_search_reads_the_base_as_it_was; mod a_secret_is_sealed_at_rest; mod a_signature_holds_on_every_path; +mod a_similar_name_is_proposed_not_merged; mod a_source_kind_is_listed_once; mod a_source_reaches_only_where_it_was_granted; mod a_statement_is_attested_by_what_its_document_says; @@ -105,6 +106,7 @@ mod an_amount_outlives_adoption; mod an_automatic_merge_is_gated_by_what_it_can_undo; mod an_axiom_that_moved_closes_its_findings; mod an_earlier_mention_keeps_the_stated_end; +mod an_empty_ontology_does_not_cover_the_corpus; mod an_end_date_closes_the_open_span; mod an_event_holds_at_the_moment_it_names; mod an_exploration_says_what_it_covered; @@ -133,6 +135,7 @@ mod graph_changes; mod history_shows_the_merge_itself; mod human_type_decisions; mod materialization_is_serial; +mod merge_direction_counts_live_facts_in_each_base; mod migration_0070_runs_under_any_search_path; mod miss_dismissal; mod negative_binding_definition_edit; @@ -149,3 +152,4 @@ mod the_floor_under_retrieval; mod the_nearest_chunk_is_found_however_it_is_reached; mod the_second_clock_can_be_rewound; mod the_world_axis_reaches_the_second; +mod tied_signatures_come_in_a_fixed_order; diff --git a/crates/utopia-store/tests/store/merge_direction_counts_live_facts_in_each_base.rs b/crates/utopia-store/tests/store/merge_direction_counts_live_facts_in_each_base.rs new file mode 100644 index 000000000..2b353e206 --- /dev/null +++ b/crates/utopia-store/tests/store/merge_direction_counts_live_facts_in_each_base.rs @@ -0,0 +1,156 @@ +//! 合并方向只看有效事实；知识库限定不能改变度数和 UUID 平局规则。 + +use sqlx::PgPool; +use utopia_store::resolution::merge_direction; +use uuid::Uuid; + +async fn fact( + pool: &PgPool, + kb: Uuid, + predicate: Uuid, + subject: Uuid, + object: Uuid, +) -> anyhow::Result { + let id = Uuid::now_v7(); + sqlx::query( + "INSERT INTO facts (id, kb_id, subject_id, predicate_id, object_id) + VALUES ($1, $2, $3, $4, $5)", + ) + .bind(id) + .bind(kb) + .bind(subject) + .bind(predicate) + .bind(object) + .execute(pool) + .await?; + Ok(id) +} + +async fn assert_direction( + pool: &PgPool, + survivor: Uuid, + merged: Uuid, + reason: &str, +) -> anyhow::Result<()> { + for (left, right) in [(survivor, merged), (merged, survivor)] { + assert_eq!( + merge_direction(pool, left, right).await?, + (survivor, merged), + "{reason}; input order: {left}, {right}" + ); + } + Ok(()) +} + +#[tokio::test] +async fn merge_direction_counts_live_facts_in_each_base() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + let (org, workspace) = (Uuid::now_v7(), Uuid::now_v7()); + let (kb, other_kb) = (Uuid::now_v7(), Uuid::now_v7()); + let (predicate, other_predicate) = (Uuid::now_v7(), Uuid::now_v7()); + let (first, second) = (Uuid::now_v7(), Uuid::now_v7()); + let (older, newer) = (first.min(second), first.max(second)); + let (peer, other, other_peer) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + + sqlx::query("INSERT INTO organizations (id, name) VALUES ($1, 'merge-degree-test')") + .bind(org) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO workspaces (id, org_id, name) VALUES ($1, $2, 'merge-degree-test')") + .bind(workspace) + .bind(org) + .execute(&pool) + .await?; + for (base, relation) in [(kb, predicate), (other_kb, other_predicate)] { + sqlx::query( + "INSERT INTO knowledge_bases (id, workspace_id, name) + VALUES ($1, $2, 'merge-degree-test')", + ) + .bind(base) + .bind(workspace) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO relation_types (id, kb_id, key, label) + VALUES ($1, $2, 'related_to', 'related to')", + ) + .bind(relation) + .bind(base) + .execute(&pool) + .await?; + } + for (id, base, name) in [ + (older, kb, "Older"), + (newer, kb, "Newer"), + (peer, kb, "Peer"), + (other, other_kb, "Other"), + (other_peer, other_kb, "Other peer"), + ] { + sqlx::query("INSERT INTO entities (id, kb_id, canonical_name) VALUES ($1, $2, $3)") + .bind(id) + .bind(base) + .bind(name) + .execute(&pool) + .await?; + } + + assert_direction(&pool, older, newer, "zero degrees use the smaller UUID").await?; + + fact(&pool, kb, predicate, older, older).await?; + fact(&pool, kb, predicate, peer, newer).await?; + let outgoing = fact(&pool, kb, predicate, newer, peer).await?; + // older 的自关联只算一条；若拆成两次计数相加，2:2 会错误地留下 older。 + assert_direction( + &pool, + newer, + older, + "incoming and outgoing facts beat one self-reference", + ) + .await?; + + sqlx::query("UPDATE facts SET invalidated_at = now() WHERE id = $1") + .bind(outgoing) + .execute(&pool) + .await?; + assert_direction( + &pool, + older, + newer, + "invalidated facts do not break the one-to-one degree tie", + ) + .await?; + + fact(&pool, kb, predicate, older, peer).await?; + assert_direction(&pool, older, newer, "the larger live degree survives").await?; + + fact(&pool, other_kb, other_predicate, other, other).await?; + fact(&pool, other_kb, other_predicate, other, other_peer).await?; + fact(&pool, other_kb, other_predicate, other_peer, other).await?; + assert_direction( + &pool, + older, + newer, + "facts in another knowledge base leave the result unchanged", + ) + .await?; + assert_direction( + &pool, + other, + older, + "each entity's degree uses its own knowledge base", + ) + .await?; + + sqlx::query("DELETE FROM knowledge_bases WHERE id = ANY($1)") + .bind(vec![kb, other_kb]) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id = $1") + .bind(org) + .execute(&pool) + .await?; + Ok(()) +} diff --git a/crates/utopia-store/tests/store/tied_signatures_come_in_a_fixed_order.rs b/crates/utopia-store/tests/store/tied_signatures_come_in_a_fixed_order.rs new file mode 100644 index 000000000..6059ee85e --- /dev/null +++ b/crates/utopia-store/tests/store/tied_signatures_come_in_a_fixed_order.rs @@ -0,0 +1,207 @@ +//! 签名与类别词的扫描一趟聚完（#1096）：多于三条的组只取最早的三条例句，引文跟着自己的 +//! 例句走；计数与短语都相同、两端的类不同的签名按类的 id 定序，每次一样。 +use sqlx::PgPool; +use utopia_store::graph::{self, FactObject}; +use utopia_store::{phrase_bindings, type_bindings}; +use uuid::Uuid; + +/// 一条签名的键与计数：短语、主语的类、宾语的类、宾语是不是字面值、陈述数 +type Shape = (String, Option, Option, bool, i64); + +#[tokio::test] +async fn tied_signatures_come_in_a_fixed_order() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + utopia_store::db::migrate(&pool).await?; + let (org, ws, kb) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + // v7 按生成先后递增：a 类的 id 小于 b 类 + let (a, b) = (Uuid::now_v7(), Uuid::now_v7()); + sqlx::raw_sql(&format!( + "INSERT INTO organizations(id,name) VALUES ('{org}','tied-signatures'); + INSERT INTO workspaces(id,org_id,name) VALUES ('{ws}','{org}','tied-signatures'); + INSERT INTO knowledge_bases(id,workspace_id,name) VALUES ('{kb}','{ws}','tied-signatures'); + INSERT INTO entity_types(id,kb_id,key,label,color,shape) VALUES + ('{a}','{kb}','a','A','#000','circle'), ('{b}','{kb}','b','B','#000','circle');" + )) + .execute(&pool) + .await?; + let result = async { + // 类别词 company 下四种写法，各 4、3、2、1 个实体：写法只留前三 + let mut subjects = Vec::new(); + for (spelling, n) in [ + ("Company", 4), + ("company", 3), + ("COMPANY", 2), + ("CoMpany", 1), + ] { + for i in 0..n { + subjects.push(entity(&pool, kb, &format!("{spelling}-{i}"), a, spelling).await?); + } + } + let target = entity(&pool, kb, "target", b, "place").await?; + + // 一组五条：例句取最早的三条，引文跟着各自的陈述（第一条没有证据，引文为空） + let mut supplies = Vec::new(); + for s in &subjects[..5] { + supplies.push(say(&pool, kb, *s, "supplies", FactObject::Entity(target)).await?); + } + let doc = Uuid::now_v7(); + sqlx::query("INSERT INTO documents(id,kb_id,filename,sha256) VALUES($1,$2,$3,$3)") + .bind(doc) + .bind(kb) + .bind(doc.to_string()) + .execute(&pool) + .await?; + for (seq, fact) in supplies.iter().enumerate().skip(1) { + let (chunk, text) = (Uuid::now_v7(), format!("q{seq}")); + sqlx::query("INSERT INTO chunks(id,kb_id,document_id,seq,text) VALUES($1,$2,$3,$4,$5)") + .bind(chunk) + .bind(kb) + .bind(doc) + .bind(seq as i32) + .bind(&text) + .execute(&pool) + .await?; + graph::add_evidence_located(&pool, *fact, chunk, Some(&text), None, Some((0, 2))) + .await?; + } + + // 三组「owns」各两条，计数与短语都相同，只差两端的类与宾语是不是字面值 + let (x, y) = ( + serde_json::json!({ "value": "x" }), + serde_json::json!({ "value": "y" }), + ); + for s in &subjects[..2] { + say(&pool, kb, target, "owns", FactObject::Entity(*s)).await?; // b → a + say(&pool, kb, *s, "owns", FactObject::Entity(target)).await?; // a → b + } + say(&pool, kb, subjects[0], "owns", FactObject::Value(&x)).await?; // a → 字面值 + say(&pool, kb, subjects[1], "owns", FactObject::Value(&y)).await?; + + let sigs = phrase_bindings::signatures(&pool, kb).await?; + let shape: Vec = sigs + .iter() + .map(|s| { + ( + s.phrase.clone(), + s.subject_type_id, + s.object_type_id, + s.object_is_value, + s.count, + ) + }) + .collect(); + anyhow::ensure!( + shape + == vec![ + ("supplies".into(), Some(a), Some(b), false, 5), + // 平手按主语的类、宾语的类（空的在后）、宾语是不是字面值 + ("owns".into(), Some(a), Some(b), false, 2), + ("owns".into(), Some(a), None, true, 2), + ("owns".into(), Some(b), Some(a), false, 2), + ], + "{shape:#?}" + ); + let supplied = &sigs[0]; + anyhow::ensure!( + supplied.examples + == vec![ + "Company-0 —supplies→ target", + "Company-1 —supplies→ target", + "Company-2 —supplies→ target", + ], + "the first three by when they were recorded: {:?}", + supplied.examples + ); + anyhow::ensure!( + supplied.quotes == vec!["", "q1", "q2"], + "each quote follows its own example: {:?}", + supplied.quotes + ); + anyhow::ensure!(sigs[2].examples == vec!["Company-0 —owns→ x", "Company-1 —owns→ y"]); + anyhow::ensure!(sigs[2].quotes == vec!["", ""]); + for _ in 0..3 { + let again: Vec<_> = phrase_bindings::signatures(&pool, kb) + .await? + .iter() + .map(|s| (s.phrase.clone(), s.subject_type_id, s.object_type_id)) + .collect(); + let first: Vec<_> = sigs + .iter() + .map(|s| (s.phrase.clone(), s.subject_type_id, s.object_type_id)) + .collect(); + anyhow::ensure!(again == first, "the order holds from call to call"); + } + + let kinds = type_bindings::signatures(&pool, kb).await?; + let company = kinds + .iter() + .find(|k| k.kind_word == "company") + .ok_or_else(|| anyhow::anyhow!("no company: {kinds:?}"))?; + anyhow::ensure!(company.count == 10); + anyhow::ensure!( + company.words == vec!["Company", "company", "COMPANY"], + "the three most frequent spellings: {:?}", + company.words + ); + anyhow::ensure!(company.examples == vec!["Company-0", "Company-1", "Company-2"]); + anyhow::ensure!( + company.phrases == vec!["supplies", "owns"], + "phrases by how often a company says them: {:?}", + company.phrases + ); + anyhow::ensure!( + kinds + .iter() + .map(|k| k.kind_word.as_str()) + .collect::>() + == vec!["company", "place"] + ); + anyhow::Ok(()) + } + .await; + sqlx::query("DELETE FROM organizations WHERE id=$1") + .bind(org) + .execute(&pool) + .await?; + pool.close().await; + result +} + +async fn entity( + pool: &PgPool, + kb: Uuid, + name: &str, + type_id: Uuid, + kind_word: &str, +) -> anyhow::Result { + let id = Uuid::now_v7(); + sqlx::query( + "INSERT INTO entities (id, kb_id, canonical_name, type_id, specific_type) + VALUES ($1, $2, $3, $4, $5)", + ) + .bind(id) + .bind(kb) + .bind(name) + .bind(type_id) + .bind(kind_word) + .execute(pool) + .await?; + Ok(id) +} + +async fn say( + pool: &PgPool, + kb: Uuid, + subject: Uuid, + phrase: &str, + object: FactObject<'_>, +) -> anyhow::Result { + Ok( + graph::insert_open_statement(pool, kb, subject, phrase, object, None, 1.0) + .await? + .0, + ) +} diff --git a/docs/decisions/0045-a-time-mention-is-resolved-against-its-document.md b/docs/decisions/0045-a-time-mention-is-resolved-against-its-document.md index d7aa43d23..7d33c0a6d 100644 --- a/docs/decisions/0045-a-time-mention-is-resolved-against-its-document.md +++ b/docs/decisions/0045-a-time-mention-is-resolved-against-its-document.md @@ -63,6 +63,8 @@ Each resolved time carries a grade: **A**, an absolute date the text writes; **B The interpretation is the model's; code does arithmetic on structured fields and checks the words against the chunk. No list of relative expressions, month names or period words is matched in code, in any language. +A timestamp is not a time word: a complete date or zoned timestamp a pushed statement supplies in `when` / `ended` (a `statements` source, 0054) is structured already, so `resolve_time` reads it in code at its stated precision and keeps the literal words as the mention (#1089). Mentions that arrive as prose keep this decision's model path. + ## Not doing - **Dates computed by the model**, in any field of the contract. diff --git a/docs/decisions/0054-a-source-may-push-statements-in-the-open-contract.md b/docs/decisions/0054-a-source-may-push-statements-in-the-open-contract.md index 14e31a288..9d66596d6 100644 --- a/docs/decisions/0054-a-source-may-push-statements-in-the-open-contract.md +++ b/docs/decisions/0054-a-source-may-push-statements-in-the-open-contract.md @@ -87,7 +87,7 @@ Content-Type: application/json ## Open questions - Whether a statement with no offsets should look any different on a Review card. Today it does not. -- Whether `when` should accept an RFC 3339 instant directly rather than time words, once the `instant` precision on the roadmap exists (0045). +- Whether `when` should accept an RFC 3339 instant directly rather than time words, once the `instant` precision on the roadmap exists (0045). Settled 2026-10-06 (#1089): `when` and `ended` stay polymorphic and keep the words verbatim as the mention; a pushed complete date or zoned timestamp (`YYYY-MM-DD`, `YYYY-MM-DDTHH[:MM[:SS]]` with `Z` or an offset) is resolved by `resolve_time` itself, at the precision it writes and without a model round trip. Words code cannot read still go to the model, or wait at grade C when no chat model is configured. A bare year or year-month is not in this cut. ## Revisions diff --git a/docs/pipeline.md b/docs/pipeline.md index 42ccef3ae..f56da32c4 100644 --- a/docs/pipeline.md +++ b/docs/pipeline.md @@ -78,6 +78,11 @@ flowchart TB ### Where this stage drops things +**Failure diagnostics are stored in English.** The extraction worker writes its +chunk-failure labels, the exhausted-retry reason and the per-attempt summary into +`documents.graph_error`, and Library shows that text as stored. Provider and parser +messages are kept verbatim. + Eleven reason codes, all recorded in `extraction_drops` and visible in the UI (one is a trace, not a drop): | Reason | When | diff --git a/migrations/0106_a_clickhouse_warehouse_can_be_mounted.sql b/migrations/0106_a_clickhouse_warehouse_can_be_mounted.sql new file mode 100644 index 000000000..46ad2d57d --- /dev/null +++ b/migrations/0106_a_clickhouse_warehouse_can_be_mounted.sql @@ -0,0 +1,5 @@ +-- ClickHouse 进注册表（#1092），走它的 HTTP 接口（`query_engine/clickhouse.rs`）。 +-- 允许的名字与 `query_engine::ENGINES` 同一张表；加一个引擎仍然只放宽这一条 CHECK（同 0021、0025）。 +ALTER TABLE data_sources DROP CONSTRAINT data_sources_engine_check; +ALTER TABLE data_sources ADD CONSTRAINT data_sources_engine_check + CHECK (engine IN ('postgres', 'mysql', 'trino', 'databricks', 'snowflake', 'clickhouse')); diff --git a/scripts/bench/README.md b/scripts/bench/README.md index 0006aa9c3..1331d6915 100644 --- a/scripts/bench/README.md +++ b/scripts/bench/README.md @@ -252,6 +252,16 @@ node scripts/bench/ask.mjs --kb --recall 8 # 只量检索：那条对 `lib.mjs` 是两个测量台共用的地基——**判等必须是同一份**，各写一份 `same()` 迟早漂移，而一旦漂移，「提议对了几条」与「答案对了几条」就不是同一把尺子量出来的。 +SQL 的空字段（psql 默认显示的 NULL、空字符串）与纯空白不参与数值判等，不能当成 `0`。 +`firstRow()` 仍取第一行各列的数；`value()` 只读第一列，第一列为空就返回空，不取后面的数。 +回归测试不需要模型服务；第二条额外验证真实 PostgreSQL 的输出，只执行 SELECT： + +```bash +node --test scripts/bench/lib.test.mjs +BENCH_TEST_PSQL="psql -X -h 127.0.0.1 -p 1517 -U utopia -tAc" BENCH_TEST_DB=utopia PGPASSWORD=utopia \ + node --test scripts/bench/lib.pg.test.mjs +``` + ## 本体的两个数（0061 决定 5） `competency.mjs` 把一个库**接受了的**能力问题按人在 chat 里问的方式问一遍（走 `/kbs/{id}/chat`， diff --git a/scripts/bench/lib.mjs b/scripts/bench/lib.mjs index 1e46c8be7..5c176a3a9 100644 --- a/scripts/bench/lib.mjs +++ b/scripts/bench/lib.mjs @@ -70,6 +70,10 @@ export const psql = (sql) => run(withDb(PSQL, APP_DB), sql); export const onDb = (db, sql) => run(withDb(PSQL, db), sql); export const num = (sql) => Number(psql(sql) || 0); +// psql 默认把 NULL 与空字符串都印成空字段；Number("") 却是 0， +// 让没有数的字段命中零值真值。两条读数路径都先排除空白，真实的 0 照常读。 +const numericField = (field) => field.trim() === "" ? undefined : Number(field); + /// 一条 SQL 第一行里的所有数字。 /// /// **模型不写单列查询。** 问「平均行金额」，它跑的是 @@ -81,7 +85,7 @@ export function firstRow(db, sql) { const out = onDb(db, `SET statement_timeout = '20s'; ${sql}`); const first = out.split("\n").map((l) => l.trim()).filter((l) => l !== "" && l !== "SET")[0]; if (first === undefined) return { empty: true }; - const ns = first.split("|").map((x) => Number(x)).filter((n) => Number.isFinite(n)); + const ns = first.split("|").map(numericField).filter((n) => Number.isFinite(n)); return { ns }; } catch (e) { return { error: String(e.stderr || e.message).split("\n").filter((l) => l.trim())[0]?.slice(0, 120) }; @@ -97,7 +101,8 @@ export function value(db, sql) { // 而 -tA 下数据行不带标签——不滤掉它，每条读到的第一行都是 `SET` const first = out.split("\n").map((l) => l.trim()).filter((l) => l !== "" && l !== "SET")[0]; if (first === undefined) return { empty: true }; - return { n: Number(String(first).split("|")[0]) }; + const n = numericField(first.split("|")[0]); + return n === undefined ? { empty: true } : { n }; } catch (e) { return { error: String(e.stderr || e.message).split("\n").filter((l) => l.trim())[0]?.slice(0, 120) }; } diff --git a/scripts/bench/lib.pg.test.mjs b/scripts/bench/lib.pg.test.mjs new file mode 100644 index 000000000..911b94526 --- /dev/null +++ b/scripts/bench/lib.pg.test.mjs @@ -0,0 +1,24 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +// BENCH_TEST_PSQL is a psql command with -tAc and connection options; BENCH_TEST_DB +// names an existing test database. These probes only SELECT, with no fixture writes. +test("real psql keeps blank SQL fields out of numeric benchmark scores", { + skip: !process.env.BENCH_TEST_PSQL, +}, async () => { + assert.ok(process.env.BENCH_TEST_DB, "BENCH_TEST_DB must name the test database"); + process.env.BENCH_PSQL = process.env.BENCH_TEST_PSQL; + const { firstRow, value, roughly } = await import("./lib.mjs"); + const db = process.env.BENCH_TEST_DB; + for (const field of ["NULL::numeric", "''::text", "' '::text"]) { + const sql = `SELECT 7, ${field}`; + const { ns } = firstRow(db, sql); + assert.deepEqual(ns, [7], sql); + assert.equal(ns.some((n) => roughly(n, 0)), false, sql); + assert.deepEqual(value(db, `SELECT ${field}, 7`), { empty: true }, field); + } + assert.deepEqual(firstRow(db, "SELECT 7, 0, -2.5, 0.25"), { ns: [7, 0, -2.5, 0.25] }); + assert.deepEqual(value(db, "SELECT 0"), { n: 0 }); + assert.deepEqual(firstRow(db, "SELECT NULL::numeric"), { empty: true }); + assert.deepEqual(firstRow(db, "SELECT 0 WHERE FALSE"), { empty: true }); +}); diff --git a/scripts/bench/lib.test.mjs b/scripts/bench/lib.test.mjs new file mode 100644 index 000000000..b9fc6659b --- /dev/null +++ b/scripts/bench/lib.test.mjs @@ -0,0 +1,86 @@ +import assert from "node:assert/strict"; +import childProcess from "node:child_process"; +import { syncBuiltinESMExports } from "node:module"; +import test from "node:test"; +import { firstRow, value, same, roughly } from "./lib.mjs"; + +// Stub only psql's stdout: the production command wrapper, parsers and comparison +// functions still run, without needing a database for these boundary cases. +function stdout(t, text) { + t.mock.method(childProcess, "execFileSync", () => text); + syncBuiltinESMExports(); + t.after(() => { + t.mock.restoreAll(); + syncBuiltinESMExports(); + }); +} + +for (const [label, output] of [ + ["trailing NULL or empty string", "SET\n7|\n"], + ["leading NULL or empty string", "SET\n|7\n"], + ["interior empty field", "SET\n7||9\n"], + ["whitespace-only field", "SET\n7| \t \n"], +]) { + test(`firstRow does not let ${label} match a zero answer`, (t) => { + stdout(t, output); + const ns = firstRow("fixture", "SELECT fixture").ns; + assert.deepEqual(ns, label === "interior empty field" ? [7, 9] : [7]); + // ask.mjs searches every returned number using roughly; mapping scores use same. + assert.equal(ns.some((n) => roughly(n, 0)), false); + assert.equal(ns.some((n) => same(n, 0)), false); + }); +} + +test("firstRow keeps real zero, negatives, decimals and numbers in later columns", (t) => { + stdout(t, "SET\ncontext|0|-7|0.25|-2.5|Infinity|NaN\n999\n"); + const { ns } = firstRow("fixture", "SELECT fixture"); + assert.deepEqual(ns, [0, -7, 0.25, -2.5]); + assert.equal(ns.some((n) => roughly(n, 0)), true); +}); + +test("a first row of empty fields contains no numbers", (t) => { + stdout(t, "SET\n| |\n"); + assert.deepEqual(firstRow("fixture", "SELECT fixture"), { ns: [] }); +}); + +for (const output of ["SET\n", "SET\n\n"]) { + test(`no rows or a single blank value stays empty (${JSON.stringify(output)})`, (t) => { + stdout(t, output); + assert.deepEqual(firstRow("fixture", "SELECT fixture"), { empty: true }); + assert.deepEqual(value("fixture", "SELECT fixture"), { empty: true }); + }); +} + +for (const [label, output] of [ + ["NULL or empty", "SET\n|7\n"], + ["whitespace-only", "SET\n \t |7\n"], +]) { + test(`value keeps a ${label} first column empty instead of selecting a later number`, (t) => { + stdout(t, output); + const got = value("fixture", "SELECT fixture"); + assert.deepEqual(got, { empty: true }); + assert.equal(same(got.n, 0), false); + assert.equal(roughly(got.n, 0), false); + }); +} + +for (const n of [0, -7, 0.25]) { + test(`value preserves a numeric first column (${n})`, (t) => { + stdout(t, `SET\n${n}|999\n`); + assert.deepEqual(value("fixture", "SELECT fixture"), { n }); + }); +} + +test("a nonnumeric first column does not become a later numeric value", (t) => { + stdout(t, "SET\ntext|0\n"); + assert.ok(Number.isNaN(value("fixture", "SELECT fixture").n)); +}); + +test("SQL errors retain their diagnostic instead of supplying a number", (t) => { + stdout(t, ""); + childProcess.execFileSync.mock.mockImplementation(() => { + throw Object.assign(new Error("command failed"), { stderr: "ERROR: bad SQL\n" }); + }); + assert.deepEqual(firstRow("fixture", "SELECT fixture"), { error: "ERROR: bad SQL" }); + assert.deepEqual(value("fixture", "SELECT fixture"), { error: "ERROR: bad SQL" }); +}); diff --git a/web/src/api.ts b/web/src/api.ts index eaf8c9a1d..39413a715 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -1289,6 +1289,8 @@ export interface QuestionReport { changed: number; changed_share: number | null; }; + /** 文档说了、本体还放不下的：没有类的类别词、没绑到属性的短语形状（代理读的就是这两样） */ + uncovered: { kind_words: number; phrases: number }; } export interface OntologyMiss { diff --git a/web/src/i18n/en.ts b/web/src/i18n/en.ts index 20394dad9..fda2c99bc 100644 --- a/web/src/i18n/en.ts +++ b/web/src/i18n/en.ts @@ -1214,7 +1214,8 @@ export const en = { "mysql://user:pass@host:3306/db (MariaDB, TiDB, OceanBase, Doris, StarRocks)\n" + "trino://user[:pass]@host:8080/catalog[/schema] (Iceberg, Delta Lake, Hive)\n" + "databricks://:TOKEN@host/sql/1.0/warehouses/ID?catalog=main\n" + - "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH", + "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH\n" + + "clickhouse://user:pass@host:8123/db (?ssl=true: https, 8443)", add: "Add data source", newTitle: "New data source", engine: "Engine", @@ -1692,8 +1693,8 @@ export const en = { importBy: (who: string, when: string) => `${who} · ${when}`, importSize: (bytes: number) => bytes < 1024 ? `${bytes} B` : `${(bytes / 1024).toFixed(0)} KB`, - misses: "Unmatched from extraction", - missesHint: "Types and predicates the extractor produced outside the ontology, to judge what the ontology still lacks.", + misses: "Not in the ontology yet", + missesHint: "What the documents say that the ontology has no place for, and what the agent proposes for it.", /* ---- 一端挂着两个以上开放值的谓词（#341） ---- 文案克制：状态一句话，后果一句话，动作在按钮上。这一档的读者要判断的是 「这条关系一次只能有一个值吗」，不是读一篇关于双时态的说明 */ @@ -1729,7 +1730,16 @@ export const en = { restore: "Restore", suggest: "Suggest with AI", suggesting: "Analyzing…", - noMisses: "No unmatched types — the ontology covers your corpus.", + noMisses: "Nothing the documents say is waiting for a place in the ontology.", + uncovered: (kindWords: number, phrases: number) => { + const what = [ + kindWords > 0 ? `${kindWords} ${kindWords === 1 ? "kind of thing" : "kinds of thing"}` : "", + phrases > 0 ? `${phrases} relation ${phrases === 1 ? "phrase" : "phrases"}` : "", + ] + .filter(Boolean) + .join(" and "); + return `The documents use ${what} the ontology has no place for. Ask the agent to propose classes and properties for them; nothing is added until you accept it.`; + }, approve: "Add", /* 映射那一档的按钮。刻意不叫 Add——它不加东西，本体里已经有了。 两个按钮都写 Add 的话，"已经有了"这件事在界面上就消失了 */ @@ -1828,7 +1838,9 @@ export const en = { /* 从前这句把「先加个类或导入 OWL 文件」说成了开始的前提，而本体本来就 从语料里长（0003，默认开）——那句话正是 #313 说的劝退点。现在只说状态， 动作留给左栏本来就有的 New class 与 Import */ - schemaEmpty: "No classes yet. Extraction adds them as documents arrive.", + schemaEmpty: + "No classes yet. Documents are read without them; classes and properties are what turn what was read into typed facts. Have the agent propose them from your documents, import an ontology, or add a class yourself.", + schemaEmptyAction: "See what the documents need", schemaFitView: "Fit view", schemaZoomIn: "Zoom in", schemaZoomOut: "Zoom out", diff --git a/web/src/i18n/zh.ts b/web/src/i18n/zh.ts index c530be005..a5dd7d16d 100644 --- a/web/src/i18n/zh.ts +++ b/web/src/i18n/zh.ts @@ -1072,7 +1072,8 @@ export const zh: Strings = { "mysql://user:pass@host:3306/db （MariaDB、TiDB、OceanBase、Doris、StarRocks）\n" + "trino://user[:pass]@host:8080/catalog[/schema] （Iceberg、Delta Lake、Hive）\n" + "databricks://:TOKEN@host/sql/1.0/warehouses/ID?catalog=main\n" + - "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH", + "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH\n" + + "clickhouse://user:pass@host:8123/db （?ssl=true 走 https，8443）", add: "添加数据源", newTitle: "登记数据源", engine: "引擎", @@ -1489,8 +1490,8 @@ export const zh: Strings = { importBy: (who: string, when: string) => `${who} · ${when}`, importSize: (bytes: number) => bytes < 1024 ? `${bytes} B` : `${(bytes / 1024).toFixed(0)} KB`, - misses: "抽取中未匹配的", - missesHint: "抽取时落在本体之外的类型和谓词，用来判断本体还缺什么。", + misses: "本体里还没有的", + missesHint: "文档说到了、本体还放不下的东西，以及代理为它们提的建议。", /* 一端挂着两个以上开放值的谓词（#341）：状态一句话，后果一句话，动作在按钮上 */ uniqueness: "并存的取值", uniquenessHint: "标明哪些关系同时只能有一个值，之后新值会自动结束旧值。", @@ -1513,7 +1514,16 @@ export const zh: Strings = { restore: "撤回忽略", suggest: "用 AI 建议", suggesting: "分析中…", - noMisses: "没有未匹配的类型——本体覆盖了你的语料。", + noMisses: "文档说到的东西里，没有在等本体给位置的。", + uncovered: (kindWords: number, phrases: number) => { + const what = [ + kindWords > 0 ? `${kindWords} 种东西` : "", + phrases > 0 ? `${phrases} 种关系说法` : "", + ] + .filter(Boolean) + .join("、"); + return `文档里有 ${what}，本体里还没有位置。可以让代理为它们提出类和属性；你采纳之前什么都不会加进本体。`; + }, approve: "加入", mapOver: "用已有的", willRemap: (n: number) => `将重新归类 ${n} 条事实`, @@ -1585,7 +1595,9 @@ export const zh: Strings = { proposalsChangedHint: "代理的提案里人表过态的，有几条被拒或改过再采纳。", /* ---- 模式图 ---- */ schemaDiagram: "模式图", - schemaEmpty: "还没有类。文档进来时会自动补上。", + schemaEmpty: + "还没有类。没有类文档也照常读；类和属性是把读到的内容变成类型化事实的那一步。可以让代理按你的文档提出来、导入一份本体，或者自己加一个类。", + schemaEmptyAction: "看看文档需要什么", schemaFitView: "归位", schemaZoomIn: "放大", schemaZoomOut: "缩小", diff --git a/web/src/pages/Library.tsx b/web/src/pages/Library.tsx index 807d53f89..cb8a50858 100644 --- a/web/src/pages/Library.tsx +++ b/web/src/pages/Library.tsx @@ -21,6 +21,7 @@ import { import { useKb, useKbId } from "../kb"; import { toast } from "../toast"; import { copyAndSay } from "../clipboard"; +import { pushEndpoint } from "../pushEndpoint"; import { Button, Checkbox, @@ -303,7 +304,7 @@ export function Library() { const [adding, setAdding] = useState(false); const [editing, setEditing] = useState(false); // api 来源密钥弹窗（随时可查看/轮换） - const [tokenReveal, setTokenReveal] = useState<{ sourceId: string } | null>(null); + const [tokenReveal, setTokenReveal] = useState<{ sourceId: string; kind: string } | null>(null); const [cleaning, setCleaning] = useState(false); const [reExtracting, setReExtracting] = useState(false); const [rebuilding, setRebuilding] = useState(false); @@ -666,7 +667,7 @@ export function Library() { onEdit={() => setEditing(true)} onCleanup={() => setCleaning(true)} onReExtract={canEdit ? () => setReExtracting(true) : undefined} - onToken={() => setTokenReveal({ sourceId: selectedSource.id })} + onToken={() => setTokenReveal({ sourceId: selectedSource.id, kind: selectedSource.kind })} /> )} @@ -793,11 +794,11 @@ export function Library() { {adding && ( { + onDone={(id, pushKind) => { setAdding(false); if (id) setSelection(id); - // api 来源建好直接打开密钥弹窗（onboarding：端点 + 密钥一步拿全） - if (id && isApi) setTokenReveal({ sourceId: id }); + // 推送来源建好直接打开密钥弹窗（onboarding：端点 + 密钥一步拿全） + if (id && pushKind) setTokenReveal({ sourceId: id, kind: pushKind }); invalidate(); }} /> @@ -806,6 +807,7 @@ export function Library() { setTokenReveal(null)} /> )} @@ -1173,10 +1175,12 @@ function ErrorModal({ function TokenModal({ kbId, sourceId, + kind, onClose, }: { kbId: string; sourceId: string; + kind: string; onClose: () => void; }) { const queryClient = useQueryClient(); @@ -1194,7 +1198,7 @@ function TokenModal({ }); const token = tokenQuery.data?.ingest_token ?? null; - const endpoint = `${location.origin}/api/v1/sources/${sourceId}/ingest`; + const endpoint = pushEndpoint(location.origin, sourceId, kind); return ( void; + /** 第二个参数：建的是推送来源时它的 kind（`api` / `statements`），别的不给 */ + onDone: (id?: string, pushKind?: string) => void; }) { const [kind, setKind] = useState("folder"); const [name, setName] = useState(""); @@ -1446,7 +1451,8 @@ function SourceModal({ ...(syncing ? schedule : { sync_interval_minutes: null, sync_cron: null }), }); }, - onSuccess: (data) => onDone(data.source.id, kind === "api" || kind === "statements"), + onSuccess: (data) => + onDone(data.source.id, kind === "api" || kind === "statements" ? kind : undefined), }); const valid = diff --git a/web/src/pages/Ontology.tsx b/web/src/pages/Ontology.tsx index 38cb3a7fc..d5fa09007 100644 --- a/web/src/pages/Ontology.tsx +++ b/web/src/pages/Ontology.tsx @@ -225,8 +225,11 @@ export function Ontology() { enabled: !!kb, }); - const refresh = () => + const refresh = () => { queryClient.invalidateQueries({ queryKey: ["ontology", kb?.id] }); + // 采纳、拒绝之后「本体还缺什么」的那两个数也变了 + queryClient.invalidateQueries({ queryKey: ["questionReport", kb?.id] }); + }; // 错误统一走全局 toast，不再用页面内嵌错误行 const onError = (e: unknown) => toast.error((e as Error).message); @@ -576,6 +579,7 @@ export function Ontology() { ) : ( setSel({ kind: "misses" })} entityTypes={entity_types} relationTypes={relation_types} rules={rules.data?.rules ?? []} @@ -2126,6 +2130,11 @@ function MissesPanel({ onChanged: () => void; onError: (e: unknown) => void; }) { + // 文档说了、本体还放不下的有多少（代理读的那两样）：空着的本体不等于「覆盖了语料」 + const uncovered = useQuery({ + queryKey: ["questionReport", kbId], + queryFn: () => api.questionReport(kbId), + }).data?.uncovered; // 默认收起：已忽略的是**背景信息**，不该跟待处理的挤在一起抢注意力 const [showDismissed, setShowDismissed] = useState(false); const [proposals, setProposals] = useState(null); @@ -2574,7 +2583,14 @@ function MissesPanel({ /> {misses.length === 0 ? ( -

{S.ontology.noMisses}

+ // 两个数还没取回来时不说话：先说一句「没有在等的」再改口，比空着更糟 + uncovered && ( +

+ {uncovered.kind_words + uncovered.phrases > 0 + ? S.ontology.uncovered(uncovered.kind_words, uncovered.phrases) + : S.ontology.noMisses} +

+ ) ) : (
{misses.map((m) => ( diff --git a/web/src/pages/OntologySchemaGraph.tsx b/web/src/pages/OntologySchemaGraph.tsx index dcbc2e189..d2e698278 100644 --- a/web/src/pages/OntologySchemaGraph.tsx +++ b/web/src/pages/OntologySchemaGraph.tsx @@ -78,6 +78,7 @@ import { Maximize2, ZoomIn, ZoomOut } from "lucide-react"; import type { BusinessRule, EntityTypeView, RelationTypeView } from "../api"; import { S } from "../i18n"; import { + Button, CanvasLoading, cn, Pill, @@ -582,7 +583,10 @@ export function OntologySchemaGraph({ selected, onSelect, loading = false, + onEmptyAction, }: { + /** 一个类都没有时那句话下面的按钮：带人去看文档说了什么、本体还缺什么 */ + onEmptyAction?: () => void; entityTypes: EntityTypeView[]; /** 业务规则：画成主类 → 结论类的一条紫弧，点它打开规则那一页 */ rules?: BusinessRule[]; @@ -1138,8 +1142,15 @@ export function OntologySchemaGraph({ ) : empty ? (
-
+
{S.ontology.schemaEmpty} + {onEmptyAction && ( +
+ +
+ )}
) : null} diff --git a/web/src/pages/Settings.tsx b/web/src/pages/Settings.tsx index 76a341f15..26dfa20f6 100644 --- a/web/src/pages/Settings.tsx +++ b/web/src/pages/Settings.tsx @@ -477,6 +477,24 @@ export function dsSpecs(): EngineSpec[] { build: (v) => `mysql://${auth(v.user, v.password)}${v.host}:${v.port || "3306"}/${v.database}`, }, + { + id: "clickhouse", + label: "ClickHouse", + // 用户与库都可以不填：不写用户就是服务器的 default，不写库就读所有非系统库。 + // 端口 8443 / 443 服务端自己走 https，表单里不用再问一次 + fields: [ + { key: "host", label: D.fHost, placeholder: "ch.internal" }, + { key: "port", label: D.fPort, placeholder: "8123" }, + { key: "database", label: D.fDatabase, optional: true }, + { key: "user", label: D.fUser, optional: true, placeholder: "default" }, + { key: "password", label: D.fPassword, optional: true, secret: true }, + ], + // 只填了密码就是 default 用户的密码：没有用户名，服务端不会把密码发出去 + build: (v) => + `clickhouse://${auth(v.user || (v.password ? "default" : ""), v.password)}` + + `${v.host}:${v.port || "8123"}` + + (v.database ? `/${enc(v.database)}` : ""), + }, { id: "trino", label: "Trino", diff --git a/web/src/pages/dsSpecs.test.ts b/web/src/pages/dsSpecs.test.ts index 7f5166ac6..0a824e228 100644 --- a/web/src/pages/dsSpecs.test.ts +++ b/web/src/pages/dsSpecs.test.ts @@ -25,6 +25,18 @@ describe("dsSpecs", () => { } }); + it("ClickHouse 的用户与库可以不填，端口默认 8123", () => { + const ch = dsSpecs().find((s) => s.id === "clickhouse"); + expect(ch).toBeDefined(); + expect(ch!.build({ host: "ch.internal" })).toBe("clickhouse://ch.internal:8123"); + expect(ch!.build({ host: "ch.internal", password: "pw" })).toBe( + "clickhouse://default:pw@ch.internal:8123", + ); + expect( + ch!.build({ host: "ch.internal", port: "8443", database: "sales", user: "ro", password: "p@ss" }), + ).toBe("clickhouse://ro:p%40ss@ch.internal:8443/sales"); + }); + it("连接串那一档原样返回去掉首尾空白的输入", () => { const raw = dsSpecs().find((s) => s.id === "raw"); expect(raw).toBeDefined(); diff --git a/web/src/pushEndpoint.test.ts b/web/src/pushEndpoint.test.ts new file mode 100644 index 000000000..b5c233b68 --- /dev/null +++ b/web/src/pushEndpoint.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from "vitest"; + +import { pushEndpoint } from "./pushEndpoint"; + +describe("pushEndpoint", () => { + it("api 来源推到 /ingest", () => { + expect(pushEndpoint("http://localhost:1516", "abc", "api")).toBe( + "http://localhost:1516/api/v1/sources/abc/ingest", + ); + }); + + it("statements 来源推到 /statements，不是 /ingest", () => { + expect(pushEndpoint("http://localhost:1516", "abc", "statements")).toBe( + "http://localhost:1516/api/v1/sources/abc/statements", + ); + }); +}); diff --git a/web/src/pushEndpoint.ts b/web/src/pushEndpoint.ts new file mode 100644 index 000000000..cfffa5bff --- /dev/null +++ b/web/src/pushEndpoint.ts @@ -0,0 +1,9 @@ +/** 推送来源的接口地址。 + * + * 两种推送来源各有各的端点：`api` 收文档走 `/ingest`，`statements` 收陈述走 + * `/statements`（0054）。服务端对走错门的一律回 404——令牌弹窗从前对两种都写 + * `/ingest`，照着弹窗推陈述的人拿到的就是一个没有解释的 Not found（#924）。 */ +export function pushEndpoint(origin: string, sourceId: string, kind: string): string { + const door = kind === "statements" ? "statements" : "ingest"; + return `${origin}/api/v1/sources/${sourceId}/${door}`; +} diff --git a/web/src/useEventReconnect.test.ts b/web/src/useEventReconnect.test.ts index 33fbf9c85..2e37fe5c4 100644 --- a/web/src/useEventReconnect.test.ts +++ b/web/src/useEventReconnect.test.ts @@ -128,9 +128,12 @@ afterEach(() => { }); describe("KB stream reconnect", () => { - it("refetches missed changes immediately for every stream key in the current KB", async () => { + it("refetches every current KB stream key and alerts after reconnect without new business events", async () => { const client = createClient(); const current = await Promise.all(STREAM_KEYS.map((head) => watch(client, [head, "kb-a", "detail"]))); + const badge = await watch(client, ["alerts", "unread"]); + const list = await watch(client, ["alerts", "list", "", 0]); + const recoveredQueries = [...current, badge, list]; const otherKb = await watch(client, ["graph", "kb-b"]); const unrelated = await watch(client, ["health"]); const inactiveKey = ["documents", "kb-a", "unmounted"]; @@ -138,20 +141,23 @@ describe("KB stream reconnect", () => { const { source } = mount(() => useKbEvents("kb-a")); expect(source.url).toBe("/api/v1/kbs/kb-a/events"); source.emit("open"); - for (const query of current) expect(query.fetch).toHaveBeenCalledTimes(1); + for (const query of recoveredQueries) expect(query.fetch).toHaveBeenCalledTimes(1); source.emit("error"); - for (const query of current) query.changeServer(); + for (const query of recoveredQueries) query.changeServer(); // staleTime 到期只让缓存变旧，不会自动请求；断线期间完全没有业务事件。 await vi.advanceTimersByTimeAsync(STREAM_STALE_MS + 1); - for (const query of current) expect(query.fetch).toHaveBeenCalledTimes(1); + for (const query of recoveredQueries) expect(query.fetch).toHaveBeenCalledTimes(1); source.emit("open"); // 重连不等普通事件的 300ms 合并窗口。 - for (const query of current) expect(query.fetch).toHaveBeenCalledTimes(2); + for (const query of recoveredQueries) expect(query.fetch).toHaveBeenCalledTimes(2); await vi.waitFor(() => { - for (const query of current) expect(query.observer.getCurrentResult().data).toBe("after disconnect"); + for (const query of recoveredQueries) expect(query.observer.getCurrentResult().data).toBe("after disconnect"); }); + source.emit("open"); + await vi.advanceTimersByTimeAsync(SETTLE_MS); + for (const query of recoveredQueries) expect(query.fetch).toHaveBeenCalledTimes(2); expect(client.getQueryState(inactiveKey)?.isInvalidated).toBe(true); expect(otherKb.fetch).toHaveBeenCalledTimes(1); expect(unrelated.fetch).toHaveBeenCalledTimes(1); @@ -173,17 +179,25 @@ describe("KB stream reconnect", () => { const client = createClient(); const graph = await watch(client, ["graph", "kb-a"]); const documents = await watch(client, ["documents", "kb-a"]); + const badge = await watch(client, ["alerts", "unread"]); + const list = await watch(client, ["alerts", "list", "", 0]); const { source } = mount(() => useKbEvents("kb-a")); source.emit("open"); source.emit("document"); source.emit("graph"); + source.emit("alert"); + source.emit("alert"); source.emit("error"); source.emit("open"); expect(graph.fetch).toHaveBeenCalledTimes(2); expect(documents.fetch).toHaveBeenCalledTimes(2); + expect(badge.fetch).toHaveBeenCalledTimes(2); + expect(list.fetch).toHaveBeenCalledTimes(2); await vi.advanceTimersByTimeAsync(SETTLE_MS); expect(graph.fetch).toHaveBeenCalledTimes(2); expect(documents.fetch).toHaveBeenCalledTimes(2); + expect(badge.fetch).toHaveBeenCalledTimes(2); + expect(list.fetch).toHaveBeenCalledTimes(2); expect(vi.getTimerCount()).toBe(0); }); @@ -209,15 +223,18 @@ describe("KB stream reconnect", () => { const client = createClient(); const oldKb = await watch(client, ["graph", "kb-a"]); const newKb = await watch(client, ["graph", "kb-b"]); + const badge = await watch(client, ["alerts", "unread"]); + const list = await watch(client, ["alerts", "list", "", 0]); const old = mount(() => useKbEvents("kb-a")); old.source.emit("open"); old.source.emit("error"); const lateOpen = [...old.source.listeners.get("open") ?? []]; const lateGraph = [...old.source.listeners.get("graph") ?? []]; + const lateAlert = [...old.source.listeners.get("alert") ?? []]; old.unmount(); const next = mount(() => useKbEvents("kb-b")); next.source.emit("open"); - for (const callback of [...lateOpen, ...lateGraph]) callback(); + for (const callback of [...lateOpen, ...lateGraph, ...lateAlert]) callback(); old.source.emit("open"); old.source.emit("document"); await vi.advanceTimersByTimeAsync(SETTLE_MS); @@ -226,10 +243,14 @@ describe("KB stream reconnect", () => { expect(vi.getTimerCount()).toBe(0); expect(oldKb.fetch).toHaveBeenCalledTimes(1); expect(newKb.fetch).toHaveBeenCalledTimes(1); + expect(badge.fetch).toHaveBeenCalledTimes(1); + expect(list.fetch).toHaveBeenCalledTimes(1); next.source.emit("error"); next.source.emit("open"); expect(oldKb.fetch).toHaveBeenCalledTimes(1); expect(newKb.fetch).toHaveBeenCalledTimes(2); + expect(badge.fetch).toHaveBeenCalledTimes(2); + expect(list.fetch).toHaveBeenCalledTimes(2); }); it("does not subscribe without a selected KB", () => {