use std::fs;
use std::path::Path;
use aether_core::events::{AgentEvent, ContextEvent, MessageEvent, StreamState, ToolEvent, TurnEvent, TurnOutcome};
use llm::{ContentBlock, ToolCallError, ToolCallRequest, ToolCallResult};
use tempfile::TempDir;
use crate::model::{SessionControlEvent, SessionEvent, SessionMeta, UserEvent};
use crate::store::SessionStore;
pub const DEFAULT_CREATED_AT: &str = "2026-01-01T00:00:00Z";
pub fn session_meta(session_id: &str, created_at: &str) -> SessionMeta {
SessionMeta {
session_id: session_id.to_string(),
cwd: "/tmp/project".into(),
model: "test-model".into(),
selected_mode: Some("planner".into()),
created_at: created_at.to_string(),
}
}
pub fn user_message(text: &str) -> SessionEvent {
user_message_with(vec![ContentBlock::text(text)])
}
pub fn user_message_with(content: Vec<ContentBlock>) -> SessionEvent {
SessionEvent::User(UserEvent::Message { content })
}
pub fn assistant_text(message_id: &str, text: &str) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Message(MessageEvent::Text {
message_id: message_id.into(),
chunk: text.into(),
is_complete: true,
}))
}
pub fn partial_text(message_id: &str, chunk: &str) -> SessionEvent {
SessionEvent::Agent(AgentEvent::text(message_id, chunk, StreamState::Partial))
}
pub fn tool_call(id: &str, name: &str, arguments: &str) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Tool(ToolEvent::Call {
request: ToolCallRequest { id: id.into(), name: name.into(), arguments: arguments.into() },
}))
}
pub fn tool_result(id: &str, name: &str, result: &str) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Tool(ToolEvent::Result {
result: ToolCallResult { id: id.into(), name: name.into(), arguments: "{}".into(), result: result.into() },
result_meta: None,
}))
}
pub fn tool_error(id: &str, name: &str, error: &str) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Tool(ToolEvent::Error {
error: ToolCallError { id: id.into(), name: name.into(), arguments: Some("{}".into()), error: error.into() },
}))
}
pub fn turn_ended(outcome: TurnOutcome) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Turn(TurnEvent::Ended { outcome }))
}
pub fn agent_switched(from: Option<&str>, to: Option<&str>) -> SessionEvent {
SessionEvent::Control(SessionControlEvent::AgentSwitched {
from: from.map(str::to_string),
to: to.map(str::to_string),
})
}
pub fn compaction_result(summary: &str, messages_removed: usize) -> SessionEvent {
SessionEvent::Agent(AgentEvent::Context(ContextEvent::CompactionResult {
summary: summary.into(),
messages_removed,
}))
}
pub struct TestStore {
directory: TempDir,
store: SessionStore,
}
impl TestStore {
pub fn new() -> Self {
let directory = TempDir::new().expect("temporary session directory");
let store = SessionStore::from_path(directory.path().to_path_buf());
Self { directory, store }
}
pub fn path(&self) -> &Path {
self.directory.path()
}
pub fn store(&self) -> &SessionStore {
&self.store
}
pub fn session(self, session_id: &str, events: &[SessionEvent]) -> Self {
self.append_meta(session_id, &session_meta(session_id, DEFAULT_CREATED_AT));
for event in events {
self.append(session_id, event);
}
self
}
pub fn append_meta(&self, session_id: &str, meta: &SessionMeta) {
self.store.append_meta(session_id, meta).expect("failed to append metadata");
}
pub fn append(&self, session_id: &str, event: &SessionEvent) {
self.store.append_event(session_id, event).expect("failed to append event");
}
pub fn write_raw(&self, name: &str, content: &str) {
fs::write(self.path().join(name), content).expect("failed to write raw store file");
}
}
impl Default for TestStore {
fn default() -> Self {
Self::new()
}
}