#![allow(clippy::unwrap_used)]
#![allow(clippy::expect_used)]
mod common;
use std::sync::Arc;
use std::time::Duration;
use chat_engine::domain::ports::StreamEventBuffer;
use chat_engine::domain::service::message_service::{MessageService, SendMessageRequest};
use chat_engine::domain::service::plugin_service::PluginService;
use chat_engine::domain::service::session_service::Identity;
use chat_engine::infra::db::repo::stream_event_repo::SeaStreamEventBuffer;
use chat_engine_sdk::models::{FileCitation, MessagePartInput, MessagePartType};
use chat_engine_sdk::{
ChatEngineBackendPlugin, PluginError, StreamingChunkEvent, StreamingCompleteEvent,
StreamingEvent, StreamingPartEvent, StreamingStateEvent, StreamingToolEvent,
};
use futures::StreamExt;
use tokio_util::sync::CancellationToken;
use toolkit::ClientHub;
use toolkit::client_hub::ClientScope;
use uuid::Uuid;
use common::db::{self, DbHarness};
use common::{FakePlugin, FakePluginScript};
const TENANT_ID: &str = "tenant-it";
const USER_ID: &str = "user-it";
fn build_service(
harness: &DbHarness,
plugin_instance_id: &str,
plugin: Arc<dyn ChatEngineBackendPlugin>,
) -> MessageService {
let hub = Arc::new(ClientHub::new());
hub.register_scoped::<dyn ChatEngineBackendPlugin>(
ClientScope::gts_id(plugin_instance_id),
plugin,
);
let plugins = PluginService::new(hub, Arc::clone(&harness.plugin_configs));
MessageService::new(
Arc::clone(&harness.sessions),
Arc::clone(&harness.session_types),
Arc::clone(&harness.messages),
plugins,
)
}
fn make_request(session_id: Uuid) -> SendMessageRequest {
SendMessageRequest {
session_id,
parts: vec![MessagePartInput {
part_type: MessagePartType::Text,
content: serde_json::json!({"text": "hello"}),
file_citations: vec![],
link_citations: vec![],
references: vec![],
}],
file_ids: vec![],
parent_message_id: None,
capabilities: None,
}
}
fn make_identity() -> Identity {
Identity::new(TENANT_ID, USER_ID, None).unwrap()
}
#[tokio::test]
async fn cancel_after_partial_chunks_persists_is_complete_false_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "cancel-persists-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let placeholder = Uuid::nil();
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::EventsThenHang(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "alpha-".into(),
}),
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "beta-".into(),
}),
]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin.clone();
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(make_request(session_id), make_identity(), cancel.clone())
.await
.expect("send_message dispatch");
let mut wire_text = String::new();
let mut chunks_seen = 0;
let mut start_seen = false;
while let Some(evt) = stream.next().await {
match evt {
StreamingEvent::Start(_) => start_seen = true,
StreamingEvent::Chunk(c) => {
wire_text.push_str(&c.chunk);
chunks_seen += 1;
if chunks_seen == 2 {
cancel.cancel();
break;
}
}
StreamingEvent::Complete(_) | StreamingEvent::Error(_) => {
panic!("driver should not have produced Complete/Error before cancel")
}
_ => {}
}
}
assert!(start_seen, "driver must emit Start before any chunk");
assert_eq!(chunks_seen, 2, "wire stream lost a chunk");
assert_eq!(wire_text, "alpha-beta-", "wire payload was tampered with");
let row = db::wait_for_finalize(&harness.db, session_id, Duration::from_secs(2)).await;
assert!(
!row.is_complete,
"cancelled assistant row MUST be is_complete=false; persisted row = {row:?}"
);
assert_eq!(
db::message_text(&harness.db, row.message_id).await,
"alpha-beta-",
"persisted partial content must equal the chunks emitted before cancel",
);
let metadata = row
.metadata
.as_ref()
.expect("finalize_assistant must write metadata on cancel");
assert_eq!(
metadata
.get("cancelled")
.and_then(serde_json::Value::as_bool),
Some(true),
"cancel finalize must stamp metadata.cancelled=true; got {metadata}",
);
assert_eq!(
metadata.get("partial").and_then(serde_json::Value::as_bool),
Some(true),
"cancel finalize must stamp metadata.partial=true; got {metadata}",
);
assert_eq!(plugin.call_count(), 1);
}
#[tokio::test]
async fn multi_part_user_message_round_trips_in_order_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "multi-part-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let plugin = FakePlugin::new(plugin_id, FakePluginScript::Events(vec![]));
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let req = SendMessageRequest {
session_id,
parts: vec![
MessagePartInput {
part_type: MessagePartType::Text,
content: serde_json::json!({"text": "look at this"}),
file_citations: vec![],
link_citations: vec![],
references: vec![],
},
MessagePartInput {
part_type: MessagePartType::Code,
content: serde_json::json!({"language": "rust", "code": "fn main() {}"}),
file_citations: vec![],
link_citations: vec![],
references: vec![],
},
MessagePartInput {
part_type: MessagePartType::Links,
content: serde_json::json!({"links": [{"url": "https://example.com"}]}),
file_citations: vec![],
link_citations: vec![],
references: vec![],
},
],
file_ids: vec![],
parent_message_id: None,
capabilities: None,
};
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(req, make_identity(), cancel)
.await
.expect("send_message dispatch");
while stream.next().await.is_some() {}
let parts = db::message_parts_ordered(&harness.db, session_id, "user").await;
let types: Vec<&str> = parts.iter().map(|(t, _, _)| t.as_str()).collect();
assert_eq!(
types,
vec!["text", "code", "links"],
"parts must persist in submitted order",
);
let numbers: Vec<i32> = parts.iter().map(|(_, n, _)| *n).collect();
assert_eq!(
numbers,
vec![0, 1, 2],
"part numbers must be 0-based and contiguous"
);
assert_eq!(
parts[1].2.get("language").and_then(|v| v.as_str()),
Some("rust"),
"code part content must round-trip verbatim",
);
}
#[tokio::test]
async fn assistant_file_citation_persists_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "citation-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let cite: FileCitation = serde_json::from_value(serde_json::json!({
"document_id": "doc-1",
"document_name": "Doc One",
"index": 1,
"quote": "the answer is 42",
"text_positions": [7],
}))
.expect("build file citation");
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::Events(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: Uuid::nil(),
chunk: "answer".into(),
}),
StreamingEvent::Complete(StreamingCompleteEvent {
message_id: Uuid::nil(),
metadata: None,
file_citations: vec![cite],
link_citations: vec![],
references: vec![],
}),
]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
.expect("send_message dispatch");
while stream.next().await.is_some() {}
let assistant = db::find_assistant_message(&harness.db, session_id)
.await
.expect("assistant row persisted");
let cites = db::file_citations_for_message(&harness.db, assistant.message_id).await;
assert_eq!(cites.len(), 1, "one file citation must be persisted");
assert_eq!(cites[0]["document_id"], "doc-1");
assert_eq!(cites[0]["index"], 1);
assert_eq!(
cites[0]["text_positions"],
serde_json::json!([7]),
"text_positions must round-trip verbatim",
);
}
#[tokio::test]
async fn send_message_stamps_tenant_and_author_against_sqlite() {
use chat_engine::infra::db::entity::message;
let harness = db::setup_sqlite().await;
let plugin_id = "tenant-stamp-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::Events(vec![StreamingEvent::Chunk(StreamingChunkEvent {
message_id: Uuid::nil(),
chunk: "ok".into(),
})]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
.expect("send_message dispatch");
while stream.next().await.is_some() {}
let rows = db::list_messages(&harness.db, session_id).await;
let user = rows
.iter()
.find(|m| matches!(m.role, message::MessageRole::User))
.expect("user message persisted");
assert_eq!(
user.tenant_id.as_deref(),
Some(TENANT_ID),
"user message must inherit the JWT tenant",
);
assert_eq!(
user.user_id.as_deref(),
Some(USER_ID),
"user message must record its JWT author",
);
let assistant = rows
.iter()
.find(|m| matches!(m.role, message::MessageRole::Assistant))
.expect("assistant stub persisted");
assert_eq!(
assistant.tenant_id.as_deref(),
Some(TENANT_ID),
"assistant message must inherit the owning tenant",
);
assert_eq!(
assistant.user_id, None,
"assistant message has no human author",
);
}
#[tokio::test]
async fn pre_stream_timeout_persists_finish_reason_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "pre-stream-timeout-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::PreError(PluginError::timeout()),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let Err(err) = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
else {
panic!("pre-stream timeout must surface as Err");
};
let dbg = format!("{err:?}");
assert!(
dbg.contains("BackendUnavailable"),
"expected BackendUnavailable, got {dbg}",
);
let row = db::find_assistant_message(&harness.db, session_id)
.await
.expect("pre-stream timeout MUST still have inserted the assistant stub");
assert!(
!row.is_complete,
"pre-stream timeout row MUST be is_complete=false; row = {row:?}",
);
assert_eq!(
db::message_text(&harness.db, row.message_id).await,
"",
"pre-stream timeout must persist empty content (no chunks observed)",
);
let metadata = row
.metadata
.as_ref()
.expect("pre-stream timeout must write metadata");
assert_eq!(
metadata.get("finish_reason").and_then(|v| v.as_str()),
Some("timeout"),
"pre-stream timeout must stamp metadata.finish_reason=timeout; got {metadata}",
);
assert_eq!(
metadata.get("partial").and_then(serde_json::Value::as_bool),
Some(true),
"errored finalize must stamp metadata.partial=true; got {metadata}",
);
}
#[tokio::test]
async fn soft_deleted_session_rejects_send_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "lifecycle-guard-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
db::force_lifecycle_state(&harness.db, session_id, "soft_deleted").await;
let plugin = FakePlugin::new(plugin_id, FakePluginScript::Events(vec![]));
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let Err(err) = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
else {
panic!("soft_deleted session must reject send_message");
};
let dbg = format!("{err:?}");
assert!(
dbg.contains("Conflict"),
"soft_deleted session must surface as Conflict, got {dbg}",
);
let rows = db::list_messages(&harness.db, session_id).await;
assert!(
rows.is_empty(),
"rejected send must not insert any messages; got {rows:?}",
);
}
#[tokio::test]
async fn delete_message_subtree_removes_whole_tree_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "subtree-delete-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let root = db::seed_message(&harness, session_id, None, 0).await;
let a = db::seed_message(&harness, session_id, Some(root), 0).await;
let b = db::seed_message(&harness, session_id, Some(root), 1).await;
let gc = db::seed_message(&harness, session_id, Some(a), 0).await;
let other = db::seed_message(&harness, session_id, None, 1).await;
let removed = harness
.messages
.delete_message_subtree(session_id, root)
.await
.expect("delete subtree");
assert_eq!(
removed, 4,
"root + a + b + gc must all be deleted; got {removed}"
);
for (label, id) in [("root", root), ("a", a), ("b", b), ("gc", gc)] {
assert!(
db::find_message(&harness.db, id).await.is_none(),
"{label} ({id}) should be gone after subtree delete",
);
}
assert!(
db::find_message(&harness.db, other).await.is_some(),
"unrelated subtree must survive the delete",
);
let again = harness
.messages
.delete_message_subtree(session_id, root)
.await
.expect("re-delete missing root");
assert_eq!(
again, 0,
"second delete of the same root must remove nothing"
);
}
#[tokio::test]
async fn cross_tenant_send_returns_not_found_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "cross-tenant-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let plugin = FakePlugin::new(plugin_id, FakePluginScript::Events(vec![]));
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let intruder = Identity::new("tenant-other", USER_ID, None).unwrap();
let cancel = CancellationToken::new();
let Err(err) = svc
.send_message(make_request(session_id), intruder, cancel)
.await
else {
panic!("cross-tenant send must reject");
};
let dbg = format!("{err:?}");
assert!(
dbg.contains("NotFound"),
"cross-tenant send must surface as NotFound (anti-enumeration), got {dbg}",
);
let rows = db::list_messages(&harness.db, session_id).await;
assert!(
rows.is_empty(),
"cross-tenant rejected send must not insert any messages; got {rows:?}",
);
}
#[tokio::test]
async fn dropped_client_stream_still_completes_generation_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "live-tail-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let placeholder = Uuid::nil();
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::Events(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "Hel".into(),
}),
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "lo".into(),
}),
StreamingEvent::Complete(StreamingCompleteEvent {
message_id: placeholder,
metadata: None,
file_citations: vec![],
link_citations: vec![],
references: vec![],
}),
]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let stream = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
.expect("send_message dispatch");
drop(stream);
let row = db::wait_for_finalize(&harness.db, session_id, Duration::from_secs(2)).await;
assert!(
row.is_complete,
"client disconnect must NOT cancel generation under live-tail; row = {row:?}",
);
assert_eq!(
db::message_text(&harness.db, row.message_id).await,
"Hello",
"full plugin output must persist even though the client left",
);
}
#[tokio::test]
async fn streamed_parts_and_metadata_persist_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "vocab-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let placeholder = Uuid::nil();
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::Events(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "Answer".into(),
}),
StreamingEvent::Part(StreamingPartEvent {
message_id: placeholder,
part: MessagePartInput {
part_type: MessagePartType::Links,
content: serde_json::json!({ "links": [{ "url": "https://example.com" }] }),
file_citations: vec![],
link_citations: vec![],
references: vec![],
},
}),
StreamingEvent::State(StreamingStateEvent {
message_id: placeholder,
state: serde_json::json!({ "phase": "final" }),
}),
StreamingEvent::Tool(StreamingToolEvent {
message_id: placeholder,
tool: "file_search".into(),
payload: serde_json::json!({ "q": "x" }),
}),
StreamingEvent::Complete(StreamingCompleteEvent {
message_id: placeholder,
metadata: Some(serde_json::json!({ "finish_reason": "stop" })),
file_citations: vec![],
link_citations: vec![],
references: vec![],
}),
]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let svc = build_service(&harness, plugin_id, plugin_dyn);
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
.expect("send_message dispatch");
while stream.next().await.is_some() {}
let row = db::wait_for_finalize(&harness.db, session_id, Duration::from_secs(2)).await;
assert!(
row.is_complete,
"completed send must finalize is_complete=true"
);
let parts = db::message_parts_ordered(&harness.db, session_id, "assistant").await;
assert_eq!(parts.len(), 2, "expected text + links parts; got {parts:?}");
assert_eq!((parts[0].0.as_str(), parts[0].1), ("text", 0));
assert_eq!((parts[1].0.as_str(), parts[1].1), ("links", 1));
let meta = row.metadata.expect("metadata present");
assert_eq!(
meta["state"]["phase"], "final",
"State event must persist under metadata.state"
);
assert_eq!(
meta["tools"][0]["tool"], "file_search",
"Tool event must persist under metadata.tools",
);
assert_eq!(
meta["finish_reason"], "stop",
"plugin metadata must be preserved"
);
}
#[tokio::test]
async fn streamed_events_are_buffered_for_resume_against_sqlite() {
let harness = db::setup_sqlite().await;
let plugin_id = "resume-buffer-plugin";
let session_type_id = db::seed_session_type(&harness, plugin_id).await;
let session_id = db::seed_active_session(&harness, TENANT_ID, USER_ID, session_type_id).await;
let buffer: Arc<dyn StreamEventBuffer> =
Arc::new(SeaStreamEventBuffer::new(Arc::clone(&harness.db)));
let placeholder = Uuid::nil();
let plugin = FakePlugin::new(
plugin_id,
FakePluginScript::Events(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "Hel".into(),
}),
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: placeholder,
chunk: "lo".into(),
}),
StreamingEvent::Complete(StreamingCompleteEvent {
message_id: placeholder,
metadata: None,
file_citations: vec![],
link_citations: vec![],
references: vec![],
}),
]),
);
let plugin_dyn: Arc<dyn ChatEngineBackendPlugin> = plugin;
let hub = Arc::new(ClientHub::new());
hub.register_scoped::<dyn ChatEngineBackendPlugin>(ClientScope::gts_id(plugin_id), plugin_dyn);
let plugins = PluginService::new(hub, Arc::clone(&harness.plugin_configs));
let svc = MessageService::new(
Arc::clone(&harness.sessions),
Arc::clone(&harness.session_types),
Arc::clone(&harness.messages),
plugins,
)
.with_stream_buffer(Arc::clone(&buffer));
let cancel = CancellationToken::new();
let mut stream = svc
.send_message(make_request(session_id), make_identity(), cancel)
.await
.expect("send_message dispatch");
while stream.next().await.is_some() {}
let row = db::wait_for_finalize(&harness.db, session_id, Duration::from_secs(2)).await;
assert!(
row.is_complete,
"completed send must finalize is_complete=true"
);
let events = buffer
.read_since(row.message_id, None)
.await
.expect("read resume buffer");
assert!(
events.len() >= 3,
"expected start + delta(s) + complete; got {events:?}",
);
let types: Vec<&str> = events
.iter()
.map(|e| e.event["type"].as_str().unwrap_or(""))
.collect();
assert_eq!(
types.first(),
Some(&"message.start"),
"first buffered event is message.start",
);
assert_eq!(
types.last(),
Some(&"message.complete"),
"last buffered event is message.complete",
);
assert!(
types.contains(&"message.text.delta"),
"text chunks must project to message.text.delta events; got {types:?}",
);
for (i, e) in events.iter().enumerate() {
assert_eq!(
e.seq, i as u64,
"buffered seq must be contiguous; got {events:?}"
);
}
}