1130 lines
41 KiB
Rust
1130 lines
41 KiB
Rust
use diesel::prelude::*;
|
|
use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
|
|
use std::collections::HashMap;
|
|
use std::fs;
|
|
use std::path::Path;
|
|
use time::OffsetDateTime;
|
|
|
|
use crate::schema::{a2ui_messages, messages, projects, sessions};
|
|
|
|
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
|
|
|
|
#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
|
|
#[diesel(table_name = projects)]
|
|
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
|
|
pub struct Project {
|
|
pub id: i32,
|
|
pub name: String,
|
|
pub path: String,
|
|
/// Sidebar state: hides this project's session rows.
|
|
pub collapsed: bool,
|
|
}
|
|
|
|
#[derive(Insertable)]
|
|
#[diesel(table_name = projects)]
|
|
struct NewProject<'a> {
|
|
name: &'a str,
|
|
path: &'a str,
|
|
}
|
|
|
|
#[derive(Associations, Clone, Debug, Identifiable, Queryable, Selectable)]
|
|
#[diesel(belongs_to(Project))]
|
|
#[diesel(table_name = sessions)]
|
|
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
|
|
pub struct Session {
|
|
pub id: i32,
|
|
pub project_id: i32,
|
|
pub title: String,
|
|
pub context_used: i32,
|
|
pub context_limit: i32,
|
|
pub last_tokens_per_second: Option<f32>,
|
|
/// Raw column value; read it through [`Session::state`].
|
|
state: String,
|
|
pub compacted_summary: Option<String>,
|
|
last_used: i64,
|
|
}
|
|
|
|
impl Session {
|
|
pub fn state(&self) -> SessionState {
|
|
SessionState::from_id(&self.state).unwrap_or_default()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub fn fixture(id: i32, project_id: i32, title: &str, state: SessionState) -> Self {
|
|
Self {
|
|
id,
|
|
project_id,
|
|
title: title.to_owned(),
|
|
context_used: 0,
|
|
context_limit: 0,
|
|
last_tokens_per_second: None,
|
|
state: state.as_id().to_owned(),
|
|
compacted_summary: None,
|
|
last_used: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Where a session sits in the sidebar. These are lifecycle states, not flags:
|
|
/// a session is exactly one of them.
|
|
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
|
pub enum SessionState {
|
|
#[default]
|
|
Normal,
|
|
Pinned,
|
|
Archived,
|
|
}
|
|
|
|
impl SessionState {
|
|
pub fn as_id(self) -> &'static str {
|
|
match self {
|
|
Self::Normal => "normal",
|
|
Self::Pinned => "pinned",
|
|
Self::Archived => "archived",
|
|
}
|
|
}
|
|
|
|
pub fn from_id(id: &str) -> Option<Self> {
|
|
match id {
|
|
"normal" => Some(Self::Normal),
|
|
"pinned" => Some(Self::Pinned),
|
|
"archived" => Some(Self::Archived),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// Sidebar grouping order: pinned, then ordinary, then archived.
|
|
pub fn rank(self) -> u8 {
|
|
match self {
|
|
Self::Pinned => 0,
|
|
Self::Normal => 1,
|
|
Self::Archived => 2,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Insertable)]
|
|
#[diesel(table_name = sessions)]
|
|
struct NewSession<'a> {
|
|
project_id: i32,
|
|
title: &'a str,
|
|
last_used: i64,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
|
|
#[diesel(table_name = messages)]
|
|
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
|
|
pub struct StoredMessage {
|
|
pub id: i32,
|
|
pub session_id: i32,
|
|
pub user: bool,
|
|
pub tool: bool,
|
|
pub reasoning: Option<String>,
|
|
pub reasoning_complete: bool,
|
|
pub content: String,
|
|
pub model_content: Option<String>,
|
|
pub system: bool,
|
|
pub compaction: bool,
|
|
pub compaction_tail_start: Option<i32>,
|
|
pub generation_duration_ms: Option<i32>,
|
|
pub input_tokens: Option<i32>,
|
|
pub cached_tokens: Option<i32>,
|
|
pub output_tokens: Option<i32>,
|
|
}
|
|
|
|
#[derive(Insertable)]
|
|
#[diesel(table_name = messages)]
|
|
struct NewMessage<'a> {
|
|
session_id: i32,
|
|
user: bool,
|
|
tool: bool,
|
|
reasoning: Option<&'a str>,
|
|
reasoning_complete: bool,
|
|
content: &'a str,
|
|
model_content: Option<&'a str>,
|
|
system: bool,
|
|
compaction: bool,
|
|
compaction_tail_start: Option<i32>,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
|
|
#[diesel(table_name = a2ui_messages)]
|
|
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
|
|
pub struct StoredA2uiMessage {
|
|
pub id: i32,
|
|
pub session_id: i32,
|
|
pub message_id: i32,
|
|
pub json: String,
|
|
pub dismissed: bool,
|
|
}
|
|
|
|
#[derive(Insertable)]
|
|
#[diesel(table_name = a2ui_messages)]
|
|
struct NewA2uiMessage<'a> {
|
|
session_id: i32,
|
|
message_id: i32,
|
|
json: &'a str,
|
|
dismissed: bool,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ProjectWithSessions {
|
|
pub project: Project,
|
|
pub sessions: Vec<Session>,
|
|
}
|
|
|
|
pub struct Database {
|
|
connection: SqliteConnection,
|
|
}
|
|
|
|
impl Database {
|
|
pub fn open(path: &Path) -> Result<Self, String> {
|
|
let parent = path
|
|
.parent()
|
|
.ok_or_else(|| "database path has no parent directory".to_owned())?;
|
|
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
|
|
let url = path
|
|
.to_str()
|
|
.ok_or_else(|| "database path is not valid UTF-8".to_owned())?;
|
|
let mut connection = SqliteConnection::establish(url).map_err(|error| error.to_string())?;
|
|
connection
|
|
.run_pending_migrations(MIGRATIONS)
|
|
.map_err(|error| error.to_string())?;
|
|
Ok(Self { connection })
|
|
}
|
|
|
|
pub fn load_projects(&mut self) -> Result<Vec<ProjectWithSessions>, String> {
|
|
let project_rows = projects::table
|
|
.order(projects::id.asc())
|
|
.select(Project::as_select())
|
|
.load(&mut self.connection)
|
|
.map_err(|error| error.to_string())?;
|
|
let session_rows = sessions::table
|
|
.order(sessions::id.asc())
|
|
.select(Session::as_select())
|
|
.load::<Session>(&mut self.connection)
|
|
.map_err(|error| error.to_string())?;
|
|
let mut sessions_by_project: HashMap<i32, Vec<Session>> = HashMap::new();
|
|
for session in session_rows {
|
|
sessions_by_project
|
|
.entry(session.project_id)
|
|
.or_default()
|
|
.push(session);
|
|
}
|
|
// Preserve sidebar groups while matching DS4's newest-used-first order.
|
|
for sessions in sessions_by_project.values_mut() {
|
|
sessions.sort_by(|left, right| {
|
|
left.state()
|
|
.rank()
|
|
.cmp(&right.state().rank())
|
|
.then_with(|| right.last_used.cmp(&left.last_used))
|
|
.then_with(|| right.id.cmp(&left.id))
|
|
});
|
|
}
|
|
|
|
Ok(project_rows
|
|
.into_iter()
|
|
.map(|project| ProjectWithSessions {
|
|
sessions: sessions_by_project.remove(&project.id).unwrap_or_default(),
|
|
project,
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
pub fn create_project(&mut self, name: &str, path: &str) -> Result<Project, String> {
|
|
diesel::insert_into(projects::table)
|
|
.values(NewProject { name, path })
|
|
.returning(Project::as_returning())
|
|
.get_result(&mut self.connection)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn set_project_collapsed(
|
|
&mut self,
|
|
project_id: i32,
|
|
collapsed: bool,
|
|
) -> Result<(), String> {
|
|
diesel::update(projects::table.find(project_id))
|
|
.set(projects::collapsed.eq(collapsed))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> {
|
|
self.connection
|
|
.transaction(|connection| {
|
|
diesel::delete(
|
|
a2ui_messages::table.filter(
|
|
a2ui_messages::session_id.eq_any(
|
|
sessions::table
|
|
.filter(sessions::project_id.eq(project_id))
|
|
.select(sessions::id),
|
|
),
|
|
),
|
|
)
|
|
.execute(connection)?;
|
|
diesel::delete(
|
|
messages::table.filter(
|
|
messages::session_id.eq_any(
|
|
sessions::table
|
|
.filter(sessions::project_id.eq(project_id))
|
|
.select(sessions::id),
|
|
),
|
|
),
|
|
)
|
|
.execute(connection)?;
|
|
diesel::delete(sessions::table.filter(sessions::project_id.eq(project_id)))
|
|
.execute(connection)?;
|
|
diesel::delete(projects::table.find(project_id)).execute(connection)?;
|
|
Ok(())
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
|
|
pub fn create_session(&mut self, project_id: i32, title: &str) -> Result<Session, String> {
|
|
diesel::insert_into(sessions::table)
|
|
.values(NewSession {
|
|
project_id,
|
|
title,
|
|
last_used: OffsetDateTime::now_utc().unix_timestamp(),
|
|
})
|
|
.returning(Session::as_returning())
|
|
.get_result(&mut self.connection)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn touch_session(&mut self, session_id: i32) -> Result<(), String> {
|
|
touch_session(&mut self.connection, session_id).map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn rename_session(&mut self, session_id: i32, title: &str) -> Result<(), String> {
|
|
let title = title.trim();
|
|
if title.is_empty() {
|
|
return Err("Session title cannot be empty.".into());
|
|
}
|
|
diesel::update(sessions::table.find(session_id))
|
|
.set(sessions::title.eq(title))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn set_session_state(
|
|
&mut self,
|
|
session_id: i32,
|
|
state: SessionState,
|
|
) -> Result<(), String> {
|
|
diesel::update(sessions::table.find(session_id))
|
|
.set(sessions::state.eq(state.as_id()))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> {
|
|
self.connection
|
|
.transaction(|connection| {
|
|
diesel::delete(
|
|
a2ui_messages::table.filter(a2ui_messages::session_id.eq(session_id)),
|
|
)
|
|
.execute(connection)?;
|
|
diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
|
|
.execute(connection)?;
|
|
diesel::delete(sessions::table.find(session_id)).execute(connection)?;
|
|
Ok(())
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
|
|
pub fn load_messages(&mut self, session_id: i32) -> Result<Vec<StoredMessage>, String> {
|
|
messages::table
|
|
.filter(messages::session_id.eq(session_id))
|
|
.order(messages::id.asc())
|
|
.select(StoredMessage::as_select())
|
|
.load(&mut self.connection)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn load_a2ui_messages(
|
|
&mut self,
|
|
session_id: i32,
|
|
) -> Result<Vec<StoredA2uiMessage>, String> {
|
|
a2ui_messages::table
|
|
.filter(a2ui_messages::session_id.eq(session_id))
|
|
.order(a2ui_messages::id.asc())
|
|
.select(StoredA2uiMessage::as_select())
|
|
.load(&mut self.connection)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn insert_a2ui_message(
|
|
&mut self,
|
|
session_id: i32,
|
|
message_id: i32,
|
|
json: &str,
|
|
) -> Result<StoredA2uiMessage, String> {
|
|
diesel::insert_into(a2ui_messages::table)
|
|
.values(NewA2uiMessage {
|
|
session_id,
|
|
message_id,
|
|
json,
|
|
dismissed: false,
|
|
})
|
|
.returning(StoredA2uiMessage::as_returning())
|
|
.get_result(&mut self.connection)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn dismiss_a2ui_surface(
|
|
&mut self,
|
|
session_id: i32,
|
|
surface_id: &str,
|
|
) -> Result<StoredMessage, String> {
|
|
let content = format!(
|
|
"A2UI surface `{surface_id}` was dismissed by the user. There is no active A2UI surface. If the user asks for UI again, create a complete new surface with createSurface and a new surfaceId; do not update any earlier surface."
|
|
);
|
|
let json = serde_json::json!({"surfaceId": surface_id}).to_string();
|
|
self.connection
|
|
.transaction(|connection| {
|
|
let message = diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content: &content,
|
|
model_content: None,
|
|
system: true,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?;
|
|
diesel::insert_into(a2ui_messages::table)
|
|
.values(NewA2uiMessage {
|
|
session_id,
|
|
message_id: message.id,
|
|
json: &json,
|
|
dismissed: true,
|
|
})
|
|
.execute(connection)?;
|
|
Ok(message)
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
|
|
pub fn update_session_context(
|
|
&mut self,
|
|
session_id: i32,
|
|
used: u32,
|
|
limit: u32,
|
|
tokens_per_second: Option<f32>,
|
|
) -> Result<(), String> {
|
|
let used = i32::try_from(used).map_err(|_| "Used context is too large to save")?;
|
|
let limit = i32::try_from(limit).map_err(|_| "Context limit is too large to save")?;
|
|
if tokens_per_second.is_some_and(|speed| !speed.is_finite() || speed < 0.0) {
|
|
return Err("Generation speed is invalid".into());
|
|
}
|
|
diesel::update(sessions::table.find(session_id))
|
|
.set((
|
|
sessions::context_used.eq(used),
|
|
sessions::context_limit.eq(limit),
|
|
sessions::last_tokens_per_second.eq(tokens_per_second),
|
|
))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn start_chat_turn(
|
|
&mut self,
|
|
session_id: i32,
|
|
prompt: &str,
|
|
model_prompt: Option<&str>,
|
|
system_messages: &[String],
|
|
reasoning: bool,
|
|
) -> Result<Vec<StoredMessage>, String> {
|
|
self.connection
|
|
.transaction(|connection| {
|
|
reactivate_archived_session(connection, session_id)?;
|
|
touch_session(connection, session_id)?;
|
|
let mut stored = Vec::with_capacity(system_messages.len() + 2);
|
|
for content in system_messages {
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content,
|
|
model_content: None,
|
|
system: true,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
}
|
|
let user = diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: true,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content: prompt,
|
|
model_content: model_prompt,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?;
|
|
let assistant = diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: reasoning.then_some(""),
|
|
reasoning_complete: !reasoning,
|
|
content: "",
|
|
model_content: None,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?;
|
|
stored.extend([user, assistant]);
|
|
Ok(stored)
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
|
|
pub fn continue_tool_turn(
|
|
&mut self,
|
|
session_id: i32,
|
|
result: &str,
|
|
queued_user: Option<&str>,
|
|
system_messages: &[String],
|
|
reasoning: bool,
|
|
) -> Result<Vec<StoredMessage>, String> {
|
|
self.connection
|
|
.transaction(|connection| {
|
|
touch_session(connection, session_id)?;
|
|
let mut stored = Vec::with_capacity(system_messages.len() + 3);
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: true,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content: result,
|
|
model_content: None,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
if let Some(content) = queued_user {
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: true,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content,
|
|
model_content: None,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
}
|
|
for content in system_messages {
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content,
|
|
model_content: None,
|
|
system: true,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
}
|
|
let assistant = diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: reasoning.then_some(""),
|
|
reasoning_complete: !reasoning,
|
|
content: "",
|
|
model_content: None,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?;
|
|
stored.push(assistant);
|
|
Ok(stored)
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
|
|
pub fn update_message(
|
|
&mut self,
|
|
id: i32,
|
|
reasoning: Option<&str>,
|
|
reasoning_complete: bool,
|
|
content: &str,
|
|
) -> Result<(), String> {
|
|
diesel::update(messages::table.find(id))
|
|
.set((
|
|
messages::reasoning.eq(reasoning),
|
|
messages::reasoning_complete.eq(reasoning_complete),
|
|
messages::content.eq(content),
|
|
))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn update_message_generation_stats(
|
|
&mut self,
|
|
id: i32,
|
|
duration_ms: u64,
|
|
input_tokens: u32,
|
|
cached_tokens: u32,
|
|
output_tokens: u32,
|
|
) -> Result<(), String> {
|
|
let duration_ms =
|
|
i32::try_from(duration_ms).map_err(|_| "Generation duration is too large to save")?;
|
|
let input_tokens =
|
|
i32::try_from(input_tokens).map_err(|_| "Input token count is too large to save")?;
|
|
let cached_tokens =
|
|
i32::try_from(cached_tokens).map_err(|_| "Cached token count is too large to save")?;
|
|
let output_tokens =
|
|
i32::try_from(output_tokens).map_err(|_| "Output token count is too large to save")?;
|
|
diesel::update(messages::table.find(id))
|
|
.set((
|
|
messages::generation_duration_ms.eq(duration_ms),
|
|
messages::input_tokens.eq(input_tokens),
|
|
messages::cached_tokens.eq(cached_tokens),
|
|
messages::output_tokens.eq(output_tokens),
|
|
))
|
|
.execute(&mut self.connection)
|
|
.map(|_| ())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
pub fn record_compaction(
|
|
&mut self,
|
|
session_id: i32,
|
|
summary: &str,
|
|
tail_start: Option<i32>,
|
|
running_jobs: Option<&str>,
|
|
context_used: u32,
|
|
context_limit: u32,
|
|
) -> Result<Vec<StoredMessage>, String> {
|
|
let context_used = i32::try_from(context_used)
|
|
.map_err(|_| "Compacted context is too large to save".to_owned())?;
|
|
let context_limit = i32::try_from(context_limit)
|
|
.map_err(|_| "Context limit is too large to save".to_owned())?;
|
|
self.connection
|
|
.transaction(|connection| {
|
|
reactivate_archived_session(connection, session_id)?;
|
|
diesel::update(sessions::table.find(session_id))
|
|
.set((
|
|
sessions::compacted_summary.eq(Some(summary)),
|
|
sessions::context_used.eq(context_used),
|
|
sessions::context_limit.eq(context_limit),
|
|
sessions::last_tokens_per_second.eq(None::<f32>),
|
|
))
|
|
.execute(connection)?;
|
|
let mut stored = Vec::with_capacity(2);
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: false,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content: summary,
|
|
model_content: None,
|
|
system: true,
|
|
compaction: true,
|
|
compaction_tail_start: tail_start,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
if let Some(content) = running_jobs {
|
|
stored.push(
|
|
diesel::insert_into(messages::table)
|
|
.values(NewMessage {
|
|
session_id,
|
|
user: false,
|
|
tool: true,
|
|
reasoning: None,
|
|
reasoning_complete: true,
|
|
content,
|
|
model_content: None,
|
|
system: false,
|
|
compaction: false,
|
|
compaction_tail_start: None,
|
|
})
|
|
.returning(StoredMessage::as_returning())
|
|
.get_result(connection)?,
|
|
);
|
|
}
|
|
Ok(stored)
|
|
})
|
|
.map_err(|error: diesel::result::Error| error.to_string())
|
|
}
|
|
}
|
|
|
|
fn touch_session(
|
|
connection: &mut SqliteConnection,
|
|
session_id: i32,
|
|
) -> Result<(), diesel::result::Error> {
|
|
diesel::update(sessions::table.find(session_id))
|
|
.set(sessions::last_used.eq(OffsetDateTime::now_utc().unix_timestamp()))
|
|
.execute(connection)
|
|
.map(|_| ())
|
|
}
|
|
|
|
fn reactivate_archived_session(
|
|
connection: &mut SqliteConnection,
|
|
session_id: i32,
|
|
) -> Result<(), diesel::result::Error> {
|
|
diesel::update(
|
|
sessions::table
|
|
.find(session_id)
|
|
.filter(sessions::state.eq(SessionState::Archived.as_id())),
|
|
)
|
|
.set(sessions::state.eq(SessionState::Normal.as_id()))
|
|
.execute(connection)
|
|
.map(|_| ())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
#[test]
|
|
fn migrations_and_crud_work() {
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!("ds4-server-{id}.sqlite3"));
|
|
let mut database = Database::open(&path).unwrap();
|
|
|
|
let project = database.create_project("DS4", "/tmp/ds4").unwrap();
|
|
let first = database
|
|
.create_session(project.id, "First session")
|
|
.unwrap();
|
|
database
|
|
.create_session(project.id, "Second session")
|
|
.unwrap();
|
|
database.delete_session(first.id).unwrap();
|
|
let loaded = database.load_projects().unwrap();
|
|
assert_eq!(loaded[0].sessions[0].title, "Second session");
|
|
assert_eq!(loaded[0].sessions[0].state(), SessionState::Normal);
|
|
|
|
let ordinary = loaded[0].sessions[0].id;
|
|
let pinned = database.create_session(project.id, "Pinned").unwrap();
|
|
let archived = database.create_session(project.id, "Archived").unwrap();
|
|
database
|
|
.set_session_state(pinned.id, SessionState::Pinned)
|
|
.unwrap();
|
|
database
|
|
.set_session_state(archived.id, SessionState::Archived)
|
|
.unwrap();
|
|
database.rename_session(pinned.id, " Renamed ").unwrap();
|
|
assert!(database.rename_session(pinned.id, " ").is_err());
|
|
let loaded = database.load_projects().unwrap();
|
|
assert_eq!(
|
|
loaded[0]
|
|
.sessions
|
|
.iter()
|
|
.map(|session| (session.id, session.state()))
|
|
.collect::<Vec<_>>(),
|
|
[
|
|
(pinned.id, SessionState::Pinned),
|
|
(ordinary, SessionState::Normal),
|
|
(archived.id, SessionState::Archived),
|
|
]
|
|
);
|
|
assert_eq!(loaded[0].sessions[0].title, "Renamed");
|
|
// Archiving a pinned session moves it out of the pinned group entirely.
|
|
database
|
|
.set_session_state(pinned.id, SessionState::Archived)
|
|
.unwrap();
|
|
let loaded = database.load_projects().unwrap();
|
|
assert_eq!(loaded[0].sessions[0].id, ordinary);
|
|
|
|
assert!(!loaded[0].project.collapsed);
|
|
database.set_project_collapsed(project.id, true).unwrap();
|
|
assert!(database.load_projects().unwrap()[0].project.collapsed);
|
|
|
|
database.delete_project(project.id).unwrap();
|
|
assert!(database.load_projects().unwrap().is_empty());
|
|
drop(database);
|
|
|
|
let mut reopened = Database::open(&path).unwrap();
|
|
assert!(reopened.load_projects().unwrap().is_empty());
|
|
drop(reopened);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn sessions_are_ordered_by_recent_use_within_sidebar_groups() {
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!("ds4-session-order-{id}.sqlite3"));
|
|
let mut database = Database::open(&path).unwrap();
|
|
let project = database
|
|
.create_project("DS4", "/tmp/ds4-session-order")
|
|
.unwrap();
|
|
let older = database.create_session(project.id, "Older").unwrap();
|
|
let newer = database.create_session(project.id, "Newer").unwrap();
|
|
let pinned = database.create_session(project.id, "Pinned").unwrap();
|
|
database
|
|
.set_session_state(pinned.id, SessionState::Pinned)
|
|
.unwrap();
|
|
diesel::update(sessions::table.find(older.id))
|
|
.set(sessions::last_used.eq(100_i64))
|
|
.execute(&mut database.connection)
|
|
.unwrap();
|
|
diesel::update(sessions::table.find(newer.id))
|
|
.set(sessions::last_used.eq(200_i64))
|
|
.execute(&mut database.connection)
|
|
.unwrap();
|
|
|
|
let loaded = database.load_projects().unwrap();
|
|
assert_eq!(
|
|
loaded[0]
|
|
.sessions
|
|
.iter()
|
|
.map(|session| session.id)
|
|
.collect::<Vec<_>>(),
|
|
[pinned.id, newer.id, older.id]
|
|
);
|
|
|
|
drop(database);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn context_work_reactivates_an_archived_session() {
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!("ds4-reactivate-{id}.sqlite3"));
|
|
let mut database = Database::open(&path).unwrap();
|
|
let project = database
|
|
.create_project("DS4", "/tmp/ds4-reactivate")
|
|
.unwrap();
|
|
let session = database.create_session(project.id, "Archived").unwrap();
|
|
database
|
|
.set_session_state(session.id, SessionState::Archived)
|
|
.unwrap();
|
|
|
|
database
|
|
.start_chat_turn(session.id, "Resume", None, &[], false)
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
database.load_projects().unwrap()[0].sessions[0].state(),
|
|
SessionState::Normal
|
|
);
|
|
database
|
|
.set_session_state(session.id, SessionState::Archived)
|
|
.unwrap();
|
|
database
|
|
.record_compaction(session.id, "Summary", None, None, 100, 1_000)
|
|
.unwrap();
|
|
assert_eq!(
|
|
database.load_projects().unwrap()[0].sessions[0].state(),
|
|
SessionState::Normal
|
|
);
|
|
drop(database);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn a2ui_dismissal_persists_a_fresh_surface_boundary() {
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!("ds4-a2ui-dismiss-{id}.sqlite3"));
|
|
let mut database = Database::open(&path).unwrap();
|
|
let project = database
|
|
.create_project("DS4", "/tmp/ds4-a2ui-dismiss")
|
|
.unwrap();
|
|
let session = database.create_session(project.id, "A2UI").unwrap();
|
|
let first = database
|
|
.start_chat_turn(session.id, "First", None, &[], false)
|
|
.unwrap()
|
|
.pop()
|
|
.unwrap();
|
|
database
|
|
.insert_a2ui_message(
|
|
session.id,
|
|
first.id,
|
|
r#"{"version":"v1.0","createSurface":{"surfaceId":"first","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
|
|
)
|
|
.unwrap();
|
|
let dismissal = database.dismiss_a2ui_surface(session.id, "first").unwrap();
|
|
assert!(dismissal.system);
|
|
assert!(dismissal.content.contains("createSurface"));
|
|
let protocol = database.load_a2ui_messages(session.id).unwrap();
|
|
let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
|
|
(
|
|
message.id,
|
|
message.message_id,
|
|
message.dismissed,
|
|
message.json.as_str(),
|
|
)
|
|
}));
|
|
assert!(errors.is_empty());
|
|
assert!(active.active_surface().is_none());
|
|
|
|
let second = database
|
|
.start_chat_turn(session.id, "Second", None, &[], false)
|
|
.unwrap()
|
|
.pop()
|
|
.unwrap();
|
|
database
|
|
.insert_a2ui_message(
|
|
session.id,
|
|
second.id,
|
|
r#"{"version":"v1.0","createSurface":{"surfaceId":"second","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
|
|
)
|
|
.unwrap();
|
|
drop(database);
|
|
|
|
let mut reopened = Database::open(&path).unwrap();
|
|
let protocol = reopened.load_a2ui_messages(session.id).unwrap();
|
|
assert_eq!(protocol.len(), 3);
|
|
assert!(protocol[1].dismissed);
|
|
assert_eq!(protocol[1].message_id, dismissal.id);
|
|
let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
|
|
(
|
|
message.id,
|
|
message.message_id,
|
|
message.dismissed,
|
|
message.json.as_str(),
|
|
)
|
|
}));
|
|
assert!(errors.is_empty());
|
|
assert_eq!(active.active_surface().unwrap().id, "second");
|
|
let messages = reopened.load_messages(session.id).unwrap();
|
|
assert!(
|
|
messages
|
|
.iter()
|
|
.position(|message| message.id == dismissal.id)
|
|
.is_some_and(|index| messages[index + 1].content == "Second")
|
|
);
|
|
reopened.delete_project(project.id).unwrap();
|
|
drop(reopened);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn chat_and_compaction_history_survive_reopen_and_session_deletion() {
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!("ds4-chat-{id}.sqlite3"));
|
|
let mut database = Database::open(&path).unwrap();
|
|
let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap();
|
|
let session = database.create_session(project.id, "Chat").unwrap();
|
|
let mut opening = database
|
|
.start_chat_turn(
|
|
session.id,
|
|
"Question",
|
|
Some("Question\n\nhidden metadata"),
|
|
&["Date context".into()],
|
|
true,
|
|
)
|
|
.unwrap();
|
|
let assistant = opening.pop().unwrap();
|
|
database
|
|
.update_message(assistant.id, Some("Reasoning"), true, "Answer")
|
|
.unwrap();
|
|
database
|
|
.update_message_generation_stats(assistant.id, 12_345, 1_024, 768, 256)
|
|
.unwrap();
|
|
database
|
|
.insert_a2ui_message(
|
|
session.id,
|
|
assistant.id,
|
|
r#"{"version":"v1.0","createSurface":{"surfaceId":"saved","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
|
|
)
|
|
.unwrap();
|
|
database
|
|
.continue_tool_turn(
|
|
session.id,
|
|
"Tool result",
|
|
Some("Queued correction"),
|
|
&["Tool reminder".into()],
|
|
false,
|
|
)
|
|
.unwrap();
|
|
database
|
|
.update_session_context(session.id, 1_234, 65_536, Some(12.5))
|
|
.unwrap();
|
|
drop(database);
|
|
|
|
let mut reopened = Database::open(&path).unwrap();
|
|
let projects = reopened.load_projects().unwrap();
|
|
assert_eq!(projects[0].sessions[0].context_used, 1_234);
|
|
assert_eq!(projects[0].sessions[0].context_limit, 65_536);
|
|
assert_eq!(projects[0].sessions[0].last_tokens_per_second, Some(12.5));
|
|
let messages = reopened.load_messages(session.id).unwrap();
|
|
assert_eq!(messages.len(), 7);
|
|
assert!(messages[0].system);
|
|
assert_eq!(messages[1].content, "Question");
|
|
assert_eq!(
|
|
messages[1].model_content.as_deref(),
|
|
Some("Question\n\nhidden metadata")
|
|
);
|
|
assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning"));
|
|
assert!(messages[2].reasoning_complete);
|
|
assert_eq!(messages[2].content, "Answer");
|
|
assert_eq!(messages[2].generation_duration_ms, Some(12_345));
|
|
assert_eq!(messages[2].input_tokens, Some(1_024));
|
|
assert_eq!(messages[2].cached_tokens, Some(768));
|
|
assert_eq!(messages[2].output_tokens, Some(256));
|
|
assert!(messages[3].tool);
|
|
assert_eq!(messages[3].content, "Tool result");
|
|
assert!(messages[4].user);
|
|
assert_eq!(messages[4].content, "Queued correction");
|
|
assert!(messages[5].system);
|
|
assert_eq!(messages[5].content, "Tool reminder");
|
|
assert!(!messages[6].user);
|
|
assert!(!messages[6].tool);
|
|
let a2ui = reopened.load_a2ui_messages(session.id).unwrap();
|
|
assert_eq!(a2ui.len(), 1);
|
|
assert_eq!(a2ui[0].message_id, messages[2].id);
|
|
assert!(a2ui[0].json.contains("createSurface"));
|
|
let first = reopened
|
|
.record_compaction(
|
|
session.id,
|
|
"First durable state.",
|
|
Some(messages[4].id),
|
|
Some("bash job=1 status=running"),
|
|
321,
|
|
32_768,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(first.len(), 2);
|
|
assert!(first[0].compaction);
|
|
assert_eq!(first[0].compaction_tail_start, Some(messages[4].id));
|
|
drop(reopened);
|
|
let mut reopened = Database::open(&path).unwrap();
|
|
assert_eq!(
|
|
reopened.load_projects().unwrap()[0].sessions[0]
|
|
.compacted_summary
|
|
.as_deref(),
|
|
Some("First durable state.")
|
|
);
|
|
assert_eq!(
|
|
reopened.load_projects().unwrap()[0].sessions[0].context_used,
|
|
321
|
|
);
|
|
let messages = reopened.load_messages(session.id).unwrap();
|
|
assert_eq!(messages.len(), 9);
|
|
assert_eq!(messages[1].content, "Question");
|
|
assert!(messages[7].compaction);
|
|
assert_eq!(messages[7].content, "First durable state.");
|
|
assert_eq!(messages[8].content, "bash job=1 status=running");
|
|
reopened
|
|
.continue_tool_turn(session.id, "Reloaded tool result", None, &[], false)
|
|
.unwrap();
|
|
let continued = reopened.load_messages(session.id).unwrap();
|
|
let second_tail = continued[9].id;
|
|
reopened
|
|
.record_compaction(
|
|
session.id,
|
|
"Second durable state.",
|
|
Some(second_tail),
|
|
None,
|
|
222,
|
|
32_768,
|
|
)
|
|
.unwrap();
|
|
reopened
|
|
.continue_tool_turn(session.id, "Final tool result", None, &[], false)
|
|
.unwrap();
|
|
let continued = reopened.load_messages(session.id).unwrap();
|
|
let third_tail = continued[12].id;
|
|
reopened
|
|
.record_compaction(
|
|
session.id,
|
|
"Third durable state.",
|
|
Some(third_tail),
|
|
None,
|
|
111,
|
|
32_768,
|
|
)
|
|
.unwrap();
|
|
reopened
|
|
.start_chat_turn(session.id, "After third compaction", None, &[], false)
|
|
.unwrap();
|
|
drop(reopened);
|
|
|
|
let mut reopened = Database::open(&path).unwrap();
|
|
let history = reopened.load_messages(session.id).unwrap();
|
|
assert_eq!(history.len(), 17);
|
|
assert_eq!(history[1].content, "Question");
|
|
let markers = history
|
|
.iter()
|
|
.filter(|message| message.compaction)
|
|
.collect::<Vec<_>>();
|
|
assert_eq!(markers.len(), 3);
|
|
assert_eq!(markers[0].content, "First durable state.");
|
|
assert_eq!(markers[1].content, "Second durable state.");
|
|
assert_eq!(markers[2].content, "Third durable state.");
|
|
assert_eq!(markers[2].compaction_tail_start, Some(third_tail));
|
|
assert_eq!(history[15].content, "After third compaction");
|
|
reopened.delete_session(session.id).unwrap();
|
|
assert!(reopened.load_messages(session.id).unwrap().is_empty());
|
|
assert!(reopened.load_a2ui_messages(session.id).unwrap().is_empty());
|
|
drop(reopened);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
}
|