hugo/DS4Server
1
0
Fork 0
Code Issues 8 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
b182e7007c495b724f550c9f5cd2b8d87bb83ad4
DS4Server/src/database.rs
Georg Bauer b182e7007c Let the sidebar divider be dragged
2026-07-26 00:12:04 +02:00

960 lines
35 KiB
Rust
Raw Blame History

use diesel::prelude::*;
use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use crate::schema::{messages, preferences, projects, sessions};
use crate::settings::{
DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, ReasoningMode,
RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences,
StreamingCacheBudget,
};
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = preferences)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct AppPreferences {
pub id: i32,
pub selected_model: String,
pub dspark_enabled: bool,
pub idle_timeout_minutes: i32,
pub context_tokens: i32,
pub max_generated_tokens: i32,
pub system_prompt: String,
pub temperature: Option<f32>,
pub top_p: Option<f32>,
pub min_p: Option<f32>,
pub seed: Option<String>,
pub reasoning_mode: String,
pub cpu_threads: Option<i32>,
pub power_percent: Option<i32>,
pub prefill_chunk: Option<i32>,
pub quality: bool,
pub warm_weights: bool,
pub mtp_draft_tokens: i32,
pub mtp_margin: f32,
pub glm_mtp: bool,
pub glm_mtp_timing: bool,
pub dspark_confidence_threshold: Option<f32>,
pub dspark_strict: bool,
pub ssd_streaming: bool,
pub ssd_streaming_cold: bool,
pub ssd_cache_experts: Option<i64>,
pub ssd_cache_gib: Option<i64>,
pub ssd_full_layers: Option<i32>,
pub ssd_preload_experts: Option<i32>,
pub directional_steering_file: Option<String>,
pub directional_steering_ffn: Option<f32>,
pub directional_steering_attn: Option<f32>,
pub simulated_used_memory_gib: Option<i64>,
pub expert_profile_path: Option<String>,
pub endpoint_port: i32,
pub endpoint_enabled: bool,
pub endpoint_cors: bool,
pub sidebar_collapsed: bool,
/// Project the app reopens on. Cleared when that project goes away.
pub last_project_id: Option<i32>,
pub sidebar_width: i32,
}
impl Default for AppPreferences {
fn default() -> Self {
Self {
id: 1,
selected_model: "deepseek-v4-flash".into(),
dspark_enabled: false,
idle_timeout_minutes: 10,
context_tokens: 32_768,
max_generated_tokens: 50_000,
system_prompt: "You are a helpful assistant".into(),
temperature: None,
top_p: None,
min_p: None,
seed: None,
reasoning_mode: "high".into(),
cpu_threads: None,
power_percent: None,
prefill_chunk: None,
quality: false,
warm_weights: false,
mtp_draft_tokens: 1,
mtp_margin: 3.0,
glm_mtp: false,
glm_mtp_timing: false,
dspark_confidence_threshold: None,
dspark_strict: false,
ssd_streaming: false,
ssd_streaming_cold: false,
ssd_cache_experts: None,
ssd_cache_gib: None,
ssd_full_layers: None,
ssd_preload_experts: None,
directional_steering_file: None,
directional_steering_ffn: None,
directional_steering_attn: None,
simulated_used_memory_gib: None,
expert_profile_path: None,
endpoint_port: 4000,
endpoint_enabled: true,
endpoint_cors: false,
sidebar_collapsed: false,
last_project_id: None,
sidebar_width: 276,
}
}
}
impl AppPreferences {
pub(crate) fn generation(&self) -> Result<GenerationPreferences, String> {
let preferences = GenerationPreferences {
context_tokens: self.context_tokens,
max_generated_tokens: self.max_generated_tokens,
system_prompt: self.system_prompt.clone(),
temperature: self.temperature,
top_p: self.top_p,
min_p: self.min_p,
seed: self
.seed
.as_deref()
.map(str::parse)
.transpose()
.map_err(|_| "Saved seed is not a valid positive whole number.".to_owned())?,
reasoning_mode: ReasoningMode::from_id(&self.reasoning_mode)
.ok_or_else(|| format!("Unsupported reasoning mode: {}", self.reasoning_mode))?,
};
preferences.validate()?;
Ok(preferences)
}
pub(crate) fn execution(&self) -> Result<ExecutionPreferences, String> {
Ok(ExecutionPreferences {
cpu_threads: self
.cpu_threads
.map(u32::try_from)
.transpose()
.map_err(|_| "Saved CPU helper threads is invalid.".to_owned())?,
power_percent: self
.power_percent
.map(u8::try_from)
.transpose()
.map_err(|_| "Saved GPU power is invalid.".to_owned())?,
prefill_chunk: self
.prefill_chunk
.map(u32::try_from)
.transpose()
.map_err(|_| "Saved prefill chunk is invalid.".to_owned())?,
quality: self.quality,
warm_weights: self.warm_weights,
})
}
pub(crate) fn speculative(&self) -> SpeculativePreferences {
SpeculativePreferences {
mtp_draft_tokens: self.mtp_draft_tokens,
mtp_margin: self.mtp_margin,
glm_mtp: self.glm_mtp,
glm_mtp_timing: self.glm_mtp_timing,
dspark_enabled: self.dspark_enabled,
dspark_confidence_threshold: self.dspark_confidence_threshold,
dspark_strict: self.dspark_strict,
}
}
pub(crate) fn runtime(&self) -> Result<RuntimePreferences, String> {
let cache = match (self.ssd_cache_experts, self.ssd_cache_gib) {
(None, None) => None,
(Some(experts), None) => Some(StreamingCacheBudget::Experts(
u32::try_from(experts).map_err(|_| "Saved SSD expert count is invalid.")?,
)),
(None, Some(gib)) => Some(StreamingCacheBudget::Gib(
u64::try_from(gib).map_err(|_| "Saved SSD cache budget is invalid.")?,
)),
(Some(_), Some(_)) => {
return Err("Saved SSD cache budget has both count and GiB values.".into());
}
};
Ok(RuntimePreferences {
execution: self.execution()?,
speculative: self.speculative(),
ssd: SsdPreferences {
enabled: self.ssd_streaming,
cold: self.ssd_streaming_cold,
cache,
full_layers: self
.ssd_full_layers
.map(u32::try_from)
.transpose()
.map_err(|_| "Saved SSD full-layer count is invalid.")?,
preload_experts: self
.ssd_preload_experts
.map(u32::try_from)
.transpose()
.map_err(|_| "Saved SSD preload count is invalid.")?,
},
steering: SteeringPreferences {
file: self.directional_steering_file.clone(),
ffn_scale: self.directional_steering_ffn,
attention_scale: self.directional_steering_attn,
},
diagnostics: DiagnosticPreferences {
simulated_used_memory_gib: self
.simulated_used_memory_gib
.map(u64::try_from)
.transpose()
.map_err(|_| "Saved simulated memory is invalid.")?,
expert_profile_path: self.expert_profile_path.clone(),
},
})
}
}
#[derive(AsChangeset)]
#[diesel(table_name = preferences)]
struct PreferenceChanges<'a> {
selected_model: &'a str,
dspark_enabled: bool,
idle_timeout_minutes: i32,
context_tokens: i32,
max_generated_tokens: i32,
system_prompt: &'a str,
temperature: Option<f32>,
top_p: Option<f32>,
min_p: Option<f32>,
seed: Option<&'a str>,
reasoning_mode: &'a str,
cpu_threads: Option<i32>,
power_percent: Option<i32>,
prefill_chunk: Option<i32>,
quality: bool,
warm_weights: bool,
mtp_draft_tokens: i32,
mtp_margin: f32,
glm_mtp: bool,
glm_mtp_timing: bool,
dspark_confidence_threshold: Option<f32>,
dspark_strict: bool,
ssd_streaming: bool,
ssd_streaming_cold: bool,
ssd_cache_experts: Option<i64>,
ssd_cache_gib: Option<i64>,
ssd_full_layers: Option<i32>,
ssd_preload_experts: Option<i32>,
directional_steering_file: Option<&'a str>,
directional_steering_ffn: Option<f32>,
directional_steering_attn: Option<f32>,
simulated_used_memory_gib: Option<i64>,
expert_profile_path: Option<&'a str>,
endpoint_port: i32,
endpoint_enabled: bool,
endpoint_cors: bool,
}
#[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,
}
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(),
}
}
}
/// 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,
}
#[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,
}
#[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,
}
#[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);
}
// Pinned first, then ordinary, then archived; creation order within each
// group. A stable sort keeps the id ordering the query established.
for sessions in sessions_by_project.values_mut() {
sessions.sort_by_key(|session| session.state().rank());
}
Ok(project_rows
.into_iter()
.map(|project| ProjectWithSessions {
sessions: sessions_by_project.remove(&project.id).unwrap_or_default(),
project,
})
.collect())
}
pub fn load_preferences(&mut self) -> Result<AppPreferences, String> {
preferences::table
.find(1)
.select(AppPreferences::as_select())
.first(&mut self.connection)
.map_err(|error| error.to_string())
}
pub fn set_sidebar_collapsed(&mut self, collapsed: bool) -> Result<(), String> {
diesel::update(preferences::table.find(1))
.set(preferences::sidebar_collapsed.eq(collapsed))
.execute(&mut self.connection)
.map(|_| ())
.map_err(|error| error.to_string())
}
pub fn set_sidebar_width(&mut self, width: i32) -> Result<(), String> {
diesel::update(preferences::table.find(1))
.set(preferences::sidebar_width.eq(width))
.execute(&mut self.connection)
.map(|_| ())
.map_err(|error| error.to_string())
}
pub fn set_last_project(&mut self, project_id: Option<i32>) -> Result<(), String> {
diesel::update(preferences::table.find(1))
.set(preferences::last_project_id.eq(project_id))
.execute(&mut self.connection)
.map(|_| ())
.map_err(|error| error.to_string())
}
#[allow(clippy::too_many_arguments)]
pub fn update_preferences(
&mut self,
selected_model: &str,
idle_timeout_minutes: i32,
endpoint_port: i32,
endpoint_enabled: bool,
endpoint_cors: bool,
generation: &GenerationPreferences,
runtime: &RuntimePreferences,
) -> Result<AppPreferences, String> {
generation.validate()?;
let model = crate::model::ModelChoice::from_id(selected_model)
.ok_or_else(|| format!("Unsupported model: {selected_model}"))?;
runtime.validate(model)?;
if !(1..=65_535).contains(&endpoint_port) {
return Err("Endpoint port must be between 1 and 65535.".into());
}
let execution = &runtime.execution;
let speculative = &runtime.speculative;
let (ssd_cache_experts, ssd_cache_gib) = match runtime.ssd.cache {
None => (None, None),
Some(StreamingCacheBudget::Experts(experts)) => (Some(i64::from(experts)), None),
Some(StreamingCacheBudget::Gib(gib)) => (
None,
Some(i64::try_from(gib).map_err(|_| "SSD cache budget is too large.")?),
),
};
let seed = generation.seed.map(|seed| seed.to_string());
diesel::update(preferences::table.find(1))
.set(PreferenceChanges {
selected_model,
dspark_enabled: speculative.dspark_enabled,
idle_timeout_minutes,
context_tokens: generation.context_tokens,
max_generated_tokens: generation.max_generated_tokens,
system_prompt: &generation.system_prompt,
temperature: generation.temperature,
top_p: generation.top_p,
min_p: generation.min_p,
seed: seed.as_deref(),
reasoning_mode: generation.reasoning_mode.id(),
cpu_threads: execution.cpu_threads.map(|value| value as i32),
power_percent: execution.power_percent.map(i32::from),
prefill_chunk: execution.prefill_chunk.map(|value| value as i32),
quality: execution.quality,
warm_weights: execution.warm_weights,
mtp_draft_tokens: speculative.mtp_draft_tokens,
mtp_margin: speculative.mtp_margin,
glm_mtp: speculative.glm_mtp,
glm_mtp_timing: speculative.glm_mtp_timing,
dspark_confidence_threshold: speculative.dspark_confidence_threshold,
dspark_strict: speculative.dspark_strict,
ssd_streaming: runtime.ssd.enabled,
ssd_streaming_cold: runtime.ssd.cold,
ssd_cache_experts,
ssd_cache_gib,
ssd_full_layers: runtime.ssd.full_layers.map(|value| value as i32),
ssd_preload_experts: runtime.ssd.preload_experts.map(|value| value as i32),
directional_steering_file: runtime.steering.file.as_deref(),
directional_steering_ffn: runtime.steering.ffn_scale,
directional_steering_attn: runtime.steering.attention_scale,
simulated_used_memory_gib: runtime
.diagnostics
.simulated_used_memory_gib
.map(|value| value as i64),
expert_profile_path: runtime.diagnostics.expert_profile_path.as_deref(),
endpoint_port,
endpoint_enabled,
endpoint_cors,
})
.returning(AppPreferences::as_returning())
.get_result(&mut self.connection)
.map_err(|error| error.to_string())
}
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(
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 })
.returning(Session::as_returning())
.get_result(&mut self.connection)
.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(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 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,
reasoning: bool,
) -> Result<(StoredMessage, StoredMessage), String> {
self.connection
.transaction(|connection| {
let user = diesel::insert_into(messages::table)
.values(NewMessage {
session_id,
user: true,
tool: false,
reasoning: None,
reasoning_complete: true,
content: prompt,
})
.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: "",
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
Ok((user, assistant))
})
.map_err(|error: diesel::result::Error| error.to_string())
}
pub fn continue_tool_turn(
&mut self,
session_id: i32,
result: &str,
reasoning: bool,
) -> Result<(StoredMessage, StoredMessage), String> {
self.connection
.transaction(|connection| {
let tool = diesel::insert_into(messages::table)
.values(NewMessage {
session_id,
user: false,
tool: true,
reasoning: None,
reasoning_complete: true,
content: result,
})
.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: "",
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
Ok((tool, assistant))
})
.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())
}
}
#[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 preferences = database.load_preferences().unwrap();
assert_eq!(preferences.selected_model, "deepseek-v4-flash");
assert!(!preferences.dspark_enabled);
assert_eq!(preferences.idle_timeout_minutes, 10);
assert_eq!(preferences.endpoint_port, 4000);
assert!(preferences.endpoint_enabled);
assert!(!preferences.endpoint_cors);
let generation = GenerationPreferences::default();
let runtime = RuntimePreferences::default();
assert!(
database
.update_preferences(
"glm-5.2",
30,
4000,
true,
false,
&generation,
&RuntimePreferences {
speculative: SpeculativePreferences {
dspark_enabled: true,
..runtime.speculative.clone()
},
..runtime.clone()
}
)
.is_err()
);
assert!(
database
.update_preferences(
"deepseek-v4-flash",
0,
4000,
true,
false,
&generation,
&runtime,
)
.is_err()
);
let generation = GenerationPreferences {
top_p: Some(0.7),
seed: Some(u64::MAX),
reasoning_mode: ReasoningMode::Max,
..generation
};
let execution = ExecutionPreferences {
cpu_threads: Some(8),
quality: true,
..runtime.execution
};
let speculative = SpeculativePreferences {
glm_mtp: true,
glm_mtp_timing: true,
mtp_draft_tokens: 4,
mtp_margin: 2.5,
..runtime.speculative
};
let runtime = RuntimePreferences {
execution,
speculative,
ssd: SsdPreferences {
enabled: true,
cache: Some(StreamingCacheBudget::Gib(64)),
full_layers: Some(0),
preload_experts: Some(32),
..SsdPreferences::default()
},
diagnostics: DiagnosticPreferences {
simulated_used_memory_gib: Some(8),
expert_profile_path: Some("/tmp/experts.json".into()),
},
..RuntimePreferences::default()
};
database
.update_preferences("glm-5.2", 30, 4567, false, true, &generation, &runtime)
.unwrap();
let preferences = database.load_preferences().unwrap();
assert!(!preferences.endpoint_enabled);
assert!(preferences.endpoint_cors);
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);
database.set_sidebar_collapsed(true).unwrap();
database.set_last_project(Some(project.id)).unwrap();
database.set_sidebar_width(320).unwrap();
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();
let preferences = reopened.load_preferences().unwrap();
assert!(preferences.sidebar_collapsed);
assert_eq!(preferences.last_project_id, Some(project.id));
assert_eq!(preferences.sidebar_width, 320);
assert_eq!(preferences.selected_model, "glm-5.2");
assert_eq!(preferences.idle_timeout_minutes, 30);
assert_eq!(preferences.endpoint_port, 4567);
assert_eq!(preferences.generation().unwrap(), generation);
assert_eq!(preferences.runtime().unwrap(), runtime);
drop(reopened);
fs::remove_file(path).unwrap();
}
#[test]
fn chat_messages_survive_reopen_and_follow_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 (_, assistant) = database
.start_chat_turn(session.id, "Question", true)
.unwrap();
database
.update_message(assistant.id, Some("Reasoning"), true, "Answer")
.unwrap();
database
.continue_tool_turn(session.id, "Tool result", 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(), 4);
assert!(messages[0].user);
assert!(!messages[0].tool);
assert_eq!(messages[0].content, "Question");
assert_eq!(messages[1].reasoning.as_deref(), Some("Reasoning"));
assert!(messages[1].reasoning_complete);
assert_eq!(messages[1].content, "Answer");
assert!(messages[2].tool);
assert_eq!(messages[2].content, "Tool result");
assert!(!messages[3].user);
assert!(!messages[3].tool);
reopened.delete_session(session.id).unwrap();
assert!(reopened.load_messages(session.id).unwrap().is_empty());
drop(reopened);
fs::remove_file(path).unwrap();
}
}
Reference in New Issue View Git Blame Copy Permalink