From 3173f118ada68ce910fe00a882beeafa1f1e9cb6 Mon Sep 17 00:00:00 2001 From: dada-yan Date: Mon, 21 Sep 2026 17:56:29 +1000 Subject: [PATCH 1/3] Publish citation sources whenever a tool adds evidence Signed-off-by: dada-yan --- .github/workflows/ci.yml | 6 + crates/utopia-server/src/api/chat.rs | 20 +- .../src/api/chat_empty_reply_tests.rs | 3 + .../src/api/chat_sources_tests.rs | 191 ++++++++++++++++++ 4 files changed, 214 insertions(+), 6 deletions(-) create mode 100644 crates/utopia-server/src/api/chat_sources_tests.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 22fc4f962..2672bf793 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -100,6 +100,12 @@ jobs: UTOPIA_DATABASE_URL: postgres://utopia:utopia@localhost:5432/utopia UTOPIA_TEST_REQUIRE_DB: "1" + - name: Chat citation publication against Postgres + run: cargo test -p utopia-server chat_empty_reply_tests::sources_tests + env: + UTOPIA_DATABASE_URL: postgres://utopia:utopia@localhost:5432/utopia + UTOPIA_TEST_REQUIRE_DB: "1" + - name: Hybrid retrieval against Postgres run: "cargo test -p utopia-server retrieval::" env: diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index 1b08111aa..170371932 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -737,6 +737,7 @@ pub async fn chat( let mut turn_text = String::new(); let mut turn_calls: Vec = Vec::new(); let mut finished = false; + let mut published_sources = 0; while let Some(item) = run.next().await { match item { @@ -800,13 +801,20 @@ pub async fn chat( } steps_acc.push(step.clone()); yield Frame::new("step", serde_json::to_string(&step).unwrap_or_default()); - if step["kind"] == "search" || step["kind"] == "docs" { - let sources = shared.sink.lock().await.sources.clone(); - yield Frame::new( - "sources", - serde_json::to_string(&sources).unwrap_or_else(|_| "[]".into()), - ); + } + // cite() only appends: document reads can add citations too, regardless + // of the UI step kind. Release the sink before yielding to subscribers. + let sources = { + let sink = shared.sink.lock().await; + if sink.sources.len() != published_sources { + published_sources = sink.sources.len(); + Some(sink.sources.clone()) + } else { + None } + }; + if let Some(sources) = sources { + yield Frame::new("sources", serde_json::to_string(&sources).unwrap_or_else(|_| "[]".into())); } exchange_acc.push(tool_result_message(tool_result.call.as_str(), &text)); } diff --git a/crates/utopia-server/src/api/chat_empty_reply_tests.rs b/crates/utopia-server/src/api/chat_empty_reply_tests.rs index 98d53fa8f..e73a94878 100644 --- a/crates/utopia-server/src/api/chat_empty_reply_tests.rs +++ b/crates/utopia-server/src/api/chat_empty_reply_tests.rs @@ -305,3 +305,6 @@ async fn a_reply_that_stays_empty_is_an_error_after_one_retry() -> anyhow::Resul ); f.cleanup().await } + +#[path = "chat_sources_tests.rs"] +mod sources_tests; diff --git a/crates/utopia-server/src/api/chat_sources_tests.rs b/crates/utopia-server/src/api/chat_sources_tests.rs new file mode 100644 index 000000000..7e4349f22 --- /dev/null +++ b/crates/utopia-server/src/api/chat_sources_tests.rs @@ -0,0 +1,191 @@ +//! Exercise citation publication through real tools, the chat producer, reattach, +//! and stored history. The final model turn waits until the snapshot is inspected. +use super::*; +use tokio::sync::Notify; + +#[derive(Clone)] +struct SourceModel { + calls: Arc>, + seen: Arc>>, + ready: Arc, + release: Arc, +} + +async fn model( + State(m): State, + Json(body): Json, +) -> impl IntoResponse { + let index = { + let mut seen = m.seen.lock().unwrap(); + let index = seen.len(); + seen.push(body); + index + }; + let stream = async_stream::stream! { + let payload = if let Some((name, args)) = m.calls.get(index) { + json!({"choices":[{"delta":{"tool_calls":[{"index":0,"id":format!("call_{index}"),"function":{"name":name,"arguments":args.to_string()}}]},"finish_reason":"tool_calls"}]}) + } else { + m.ready.notify_one(); + m.release.notified().await; + json!({"choices":[{"delta":{"content":"The answer is in the second section [2]."},"finish_reason":"stop"}]}) + }; + yield Ok::<_, Infallible>(Event::default().data(payload.to_string())); + yield Ok::<_, Infallible>(Event::default().data("[DONE]")); + }; + Sse::new(stream) +} + +fn frames(sse: &[u8], event: &str) -> Vec { + String::from_utf8_lossy(sse) + .split("\n\n") + .filter_map(|frame| { + if !frame.lines().any(|l| l == format!("event: {event}")) { + return None; + } + frame + .lines() + .find_map(|l| l.strip_prefix("data: ")) + .and_then(|s| serde_json::from_str(s).ok()) + }) + .collect() +} + +async fn exercise(search_first: bool) -> anyhow::Result<()> { + let Some(f) = fixture(Scripted::new(vec![])).await? else { + return Ok(()); + }; + let doc = utopia_store::documents::create( + &f.pool, + f.kb, + "citation.md", + "text/plain", + 30, + "citation-test", + None, + None, + None, + ) + .await?; + let pieces: Vec<_> = ["orchard introduction", "second section contains the answer"] + .into_iter() + .enumerate() + .map(|(seq, text)| utopia_ingest::ChunkPiece { + seq: seq as i32, + text: text.into(), + char_start: 0, + char_end: text.len() as i32, + heading: None, + provenance: utopia_ingest::Provenance::stated(), + }) + .collect(); + let chunks = utopia_store::documents::replace_chunks(&f.pool, f.kb, doc.id, &pieces).await?; + f.state + .search + .reindex_document(&f.kb.to_string(), &doc.id.to_string(), &chunks)?; + let mut calls = Vec::new(); + if search_first { + calls.push(("search_chunks", json!({"query":"orchard"}))); + } + calls.extend([ + ("get_document", json!({"document_id":doc.id})), + ("get_document", json!({"document_id":doc.id})), + ("find_entities", json!({"name":"no such entity"})), + ]); + let model = SourceModel { + calls: Arc::new(calls), + seen: Default::default(), + ready: Default::default(), + release: Default::default(), + }; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let base = format!("http://{}", listener.local_addr()?); + let app = axum::Router::new() + .route("/chat/completions", axum::routing::post(self::model)) + .with_state(model.clone()); + let server = tokio::spawn(async move { axum::serve(listener, app).await }); + sqlx::query("UPDATE llm_settings SET chat_base_url=$1 WHERE workspace_id=(SELECT workspace_id FROM knowledge_bases WHERE id=$2)").bind(base).bind(f.kb).execute(&f.pool).await?; + let result = tokio::time::timeout(std::time::Duration::from_secs(20), async { + let response = chat( + State(f.state.clone()), + AuthUser(f.user.clone()), + Path(f.kb), + Json(ChatReq { + conversation_id: None, + message: "Read the second section".into(), + }), + ) + .await + .map_err(|_| anyhow::anyhow!("chat refused"))?; + model.ready.notified().await; + let id: Uuid = sqlx::query_scalar("SELECT id FROM conversations WHERE kb_id=$1") + .bind(f.kb) + .fetch_one(&f.pool) + .await?; + let snapshot = f.state.live.attach(id).await.unwrap().0; + let attached = reattach( + State(f.state.clone()), + AuthUser(f.user.clone()), + Path((f.kb, id)), + ) + .await + .map_err(|_| anyhow::anyhow!("reattach refused"))?; + model.release.notify_one(); + let live = axum::body::to_bytes(response.into_response().into_body(), 65536).await?; + let replay = axum::body::to_bytes(attached.into_response().into_body(), 65536).await?; + let Json(history) = conversation_detail( + State(f.state.clone()), + AuthUser(f.user.clone()), + Path((f.kb, id)), + ) + .await + .map_err(|_| anyhow::anyhow!("history refused"))?; + let stored = &history["messages"] + .as_array() + .unwrap() + .iter() + .find(|m| m["role"] == "assistant") + .unwrap()["sources"]; + anyhow::ensure!( + snapshot.sources.len() == 2, + "document sources missing before final answer: {:?}", + snapshot.sources + ); + let events = frames(&live, "sources"); + anyhow::ensure!( + events.len() == if search_first { 2 } else { 1 }, + "only source changes should publish: {events:?}" + ); + anyhow::ensure!(events.last() == Some(stored)); + anyhow::ensure!(json!(snapshot.sources) == *stored); + anyhow::ensure!(frames(&replay, "snapshot")[0]["sources"] == *stored); + for (index, (chunk, _)) in chunks.iter().enumerate() { + anyhow::ensure!(stored[index]["n"] == index + 1); + anyhow::ensure!(stored[index]["chunk_id"] == *chunk); + anyhow::ensure!(stored[index]["document_id"] == doc.id.to_string()); + } + let seen = model.seen.lock().unwrap(); + let final_messages = seen.last().unwrap()["messages"].to_string(); + anyhow::ensure!( + final_messages.contains("[2]") + && final_messages.contains("second section contains the answer") + ); + anyhow::ensure!(seen.len() == model.calls.len() + 1); + Ok::<_, anyhow::Error>(()) + }) + .await; + model.release.notify_one(); + server.abort(); + let _ = server.await; + f.cleanup().await?; + result??; + Ok(()) +} + +#[tokio::test] +async fn document_reads_publish_new_sources_after_search() -> anyhow::Result<()> { + exercise(true).await +} +#[tokio::test] +async fn direct_document_reads_publish_sources_without_a_search_step() -> anyhow::Result<()> { + exercise(false).await +} From f64b67dc35fb34eedb84d133a8d6de55713c66aa Mon Sep 17 00:00:00 2001 From: dada-yan Date: Mon, 21 Sep 2026 18:33:50 +1000 Subject: [PATCH 2/3] Check source publication before the final persisted replay Signed-off-by: dada-yan --- crates/utopia-server/src/api/chat_sources_tests.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/utopia-server/src/api/chat_sources_tests.rs b/crates/utopia-server/src/api/chat_sources_tests.rs index 7e4349f22..f2c2f31bd 100644 --- a/crates/utopia-server/src/api/chat_sources_tests.rs +++ b/crates/utopia-server/src/api/chat_sources_tests.rs @@ -150,7 +150,13 @@ async fn exercise(search_first: bool) -> anyhow::Result<()> { "document sources missing before final answer: {:?}", snapshot.sources ); - let events = frames(&live, "sources"); + // Count publication during tool execution, independently of a final + // persisted-source replay (e.g. the finalization work in #845). + // The last graph step adds no sources, so all citation-changing and + // duplicate document reads are before this boundary. + let live_text = String::from_utf8_lossy(&live); + let last_step = live_text.rfind("event: step\n").expect("graph step"); + let events = frames(&live_text.as_bytes()[..last_step], "sources"); anyhow::ensure!( events.len() == if search_first { 2 } else { 1 }, "only source changes should publish: {events:?}" From 4d8204e016d79f97de4c46b3f0b2f09749d73132 Mon Sep 17 00:00:00 2001 From: dada-yan Date: Mon, 21 Sep 2026 19:19:19 +1000 Subject: [PATCH 3/3] Use the shared PostgreSQL chat test step from PR 845 Signed-off-by: dada-yan --- .github/workflows/ci.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2672bf793..22fc4f962 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -100,12 +100,6 @@ jobs: UTOPIA_DATABASE_URL: postgres://utopia:utopia@localhost:5432/utopia UTOPIA_TEST_REQUIRE_DB: "1" - - name: Chat citation publication against Postgres - run: cargo test -p utopia-server chat_empty_reply_tests::sources_tests - env: - UTOPIA_DATABASE_URL: postgres://utopia:utopia@localhost:5432/utopia - UTOPIA_TEST_REQUIRE_DB: "1" - - name: Hybrid retrieval against Postgres run: "cargo test -p utopia-server retrieval::" env: