use super::*;
use crate::domain::ports::NewSession;
use crate::domain::ports::NewSessionType;
use crate::domain::session::Session;
use crate::domain::session::SessionType;
use async_trait::async_trait;
use chat_engine_sdk::error::PluginError;
use chat_engine_sdk::models::LifecycleState;
use chat_engine_sdk::plugin::{
ChatEngineBackendPlugin, MessagePluginCtx, PluginStream, stream_from_events,
};
use parking_lot::Mutex;
use std::sync::atomic::AtomicUsize;
use time::OffsetDateTime;
use toolkit::ClientHub;
use toolkit::client_hub::ClientScope;
use uuid::Uuid;
use crate::domain::message::{
Message, StreamingChunkEvent, StreamingCompleteEvent, StreamingEvent, StreamingStartEvent,
};
use crate::domain::ports::PluginConfigRepo;
use crate::domain::ports::SessionRepo;
use crate::domain::ports::SessionTypeRepo;
use crate::domain::ports::{FinalizeOutcome, InsertedPair, MessageRepo, NewUserMessage};
use crate::domain::service::PluginService;
struct MockSessionRepo {
s: Mutex<Session>,
}
impl MockSessionRepo {
fn new(session_type_id: Uuid) -> Arc<Self> {
let now = OffsetDateTime::now_utc();
Arc::new(Self {
s: Mutex::new(Session {
session_id: Uuid::new_v4(),
tenant_id: "t".into(),
user_id: "u".into(),
client_id: None,
session_type_id: Some(session_type_id),
enabled_capabilities: None,
metadata: None,
lifecycle_state: LifecycleState::Active,
share_token: None,
created_at: now,
updated_at: now,
}),
})
}
}
#[async_trait]
impl SessionRepo for MockSessionRepo {
async fn insert(&self, _m: NewSession) -> std::result::Result<Session, ChatEngineError> {
Ok(self.s.lock().clone())
}
async fn find_by_id(
&self,
t: &str,
u: &str,
id: Uuid,
) -> std::result::Result<Option<Session>, ChatEngineError> {
let s = self.s.lock().clone();
if s.tenant_id.as_str() == t && s.user_id.as_str() == u && s.session_id == id {
Ok(Some(s))
} else {
Ok(None)
}
}
async fn list_paginated(
&self,
_tenant_id: &str,
_user_id: &str,
_query: &toolkit_odata::ODataQuery,
) -> std::result::Result<toolkit_odata::Page<Session>, ChatEngineError> {
Ok(toolkit_odata::Page::empty(0))
}
async fn update_metadata(
&self,
_t: &str,
_u: &str,
_i: Uuid,
_m: Option<JsonValue>,
) -> std::result::Result<Session, ChatEngineError> {
Ok(self.s.lock().clone())
}
async fn update_capabilities(
&self,
_t: &str,
_u: &str,
_i: Uuid,
_c: Option<JsonValue>,
) -> std::result::Result<Session, ChatEngineError> {
Ok(self.s.lock().clone())
}
async fn update_lifecycle_state(
&self,
_t: &str,
_u: &str,
_i: Uuid,
_s: LifecycleState,
) -> std::result::Result<Session, ChatEngineError> {
Ok(self.s.lock().clone())
}
async fn soft_delete(
&self,
_t: &str,
_u: &str,
_i: Uuid,
_d: i64,
) -> std::result::Result<Session, ChatEngineError> {
Ok(self.s.lock().clone())
}
async fn hard_delete(
&self,
_t: &str,
_u: &str,
_i: Uuid,
) -> std::result::Result<bool, ChatEngineError> {
Ok(true)
}
}
struct MockSessionTypeRepo {
st: Mutex<SessionType>,
}
impl MockSessionTypeRepo {
fn new(id: Uuid, plugin_id: String) -> Arc<Self> {
let now = OffsetDateTime::now_utc();
Arc::new(Self {
st: Mutex::new(SessionType {
session_type_id: id,
name: "t".into(),
plugin_instance_id: Some(plugin_id),
created_at: now,
updated_at: now,
}),
})
}
}
#[async_trait]
impl SessionTypeRepo for MockSessionTypeRepo {
async fn insert(
&self,
_m: NewSessionType,
) -> std::result::Result<SessionType, ChatEngineError> {
Ok(self.st.lock().clone())
}
async fn find_by_id(
&self,
id: Uuid,
) -> std::result::Result<Option<SessionType>, ChatEngineError> {
let s = self.st.lock().clone();
if s.session_type_id == id {
Ok(Some(s))
} else {
Ok(None)
}
}
async fn list(&self) -> std::result::Result<Vec<SessionType>, ChatEngineError> {
Ok(vec![self.st.lock().clone()])
}
}
#[derive(Default)]
struct MockMessageRepo {
finalize_count: AtomicUsize,
}
#[async_trait]
impl MessageRepo for MockMessageRepo {
async fn insert_user_and_assistant_stub(
&self,
_req: NewUserMessage,
) -> std::result::Result<InsertedPair, ChatEngineError> {
Ok(InsertedPair {
user_message_id: Uuid::new_v4(),
assistant_message_id: Uuid::new_v4(),
user_variant_index: 0,
})
}
async fn finalize_assistant(
&self,
_session_id: Uuid,
_id: Uuid,
_outcome: FinalizeOutcome,
) -> std::result::Result<(), ChatEngineError> {
self.finalize_count
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(())
}
async fn fetch_active_history(
&self,
_id: Uuid,
_d: Option<u32>,
) -> std::result::Result<Vec<Message>, ChatEngineError> {
Ok(vec![])
}
async fn find_message_in_session(
&self,
_s: Uuid,
_m: Uuid,
) -> std::result::Result<Option<Message>, ChatEngineError> {
Ok(None)
}
}
struct StubPluginConfigRepo;
#[async_trait]
impl PluginConfigRepo for StubPluginConfigRepo {
async fn find(
&self,
_p: &str,
_s: Uuid,
) -> std::result::Result<Option<JsonValue>, ChatEngineError> {
Ok(None)
}
async fn upsert(
&self,
_p: &str,
_s: Uuid,
_c: JsonValue,
) -> std::result::Result<(), ChatEngineError> {
Ok(())
}
async fn delete(&self, _p: &str, _s: Uuid) -> std::result::Result<(), ChatEngineError> {
Ok(())
}
}
struct EchoPlugin {
id: String,
}
#[async_trait]
impl ChatEngineBackendPlugin for EchoPlugin {
async fn on_message(
&self,
_ctx: MessagePluginCtx,
) -> std::result::Result<PluginStream, PluginError> {
Ok(stream_from_events(vec![
StreamingEvent::Chunk(StreamingChunkEvent {
message_id: Uuid::nil(),
chunk: "hi".into(),
}),
StreamingEvent::Complete(StreamingCompleteEvent {
message_id: Uuid::nil(),
metadata: None,
file_citations: vec![],
link_citations: vec![],
references: vec![],
}),
]))
}
fn plugin_instance_id(&self) -> &str {
&self.id
}
}
fn make_service() -> Arc<MessageService> {
let session_type_id = Uuid::new_v4();
let plugin_id = "echo";
let sessions = MockSessionRepo::new(session_type_id);
let session_types = MockSessionTypeRepo::new(session_type_id, plugin_id.into());
let messages = Arc::new(MockMessageRepo::default());
let hub = Arc::new(ClientHub::new());
let plugin: Arc<dyn ChatEngineBackendPlugin> = Arc::new(EchoPlugin {
id: plugin_id.into(),
});
hub.register_scoped::<dyn ChatEngineBackendPlugin>(ClientScope::gts_id(plugin_id), plugin);
let plugin_service = PluginService::new(hub, Arc::new(StubPluginConfigRepo));
Arc::new(MessageService::new(
sessions as Arc<dyn SessionRepo>,
session_types as Arc<dyn SessionTypeRepo>,
messages as Arc<dyn MessageRepo>,
plugin_service,
))
}
#[tokio::test]
async fn ndjson_lines_serialize_one_per_event() {
let evt = StreamingEvent::Chunk(StreamingChunkEvent {
message_id: Uuid::nil(),
chunk: "hello".into(),
});
let line = serde_json::to_string(&evt).unwrap();
assert!(line.contains("\"type\":\"chunk\""));
assert!(line.contains("\"chunk\":\"hello\""));
assert!(!line.contains('\n'));
}
#[tokio::test]
async fn send_message_handler_returns_ndjson_content_type() {
let svc = make_service();
let session_id = {
Uuid::nil()
};
let _ = svc; let _ = session_id;
let evt = StreamingEvent::Start(StreamingStartEvent {
message_id: Uuid::nil(),
});
let mut buf = serde_json::to_vec(&evt).unwrap();
buf.push(b'\n');
assert!(buf.ends_with(b"\n"));
let parsed: serde_json::Value = serde_json::from_slice(&buf[..buf.len() - 1]).unwrap();
assert_eq!(parsed["type"], "start");
}