Implement session struct (#40)

MituuZ · web-flow · commit 0464725d73d3 · 2026-01-25T17:28:05.000+02:00
diff --git a/src/api/mod.rs b/src/api/mod.rs
@@ -78,6 +78,13 @@ pub fn get_chat_client_implementation(
     system_prompt: String,
     max_tokens: usize,
 ) -> Box<dyn ChatClient> {
+    // Test-only provider that returns a lightweight stub client without external dependencies
+    #[cfg(test)]
+    {
+        if provider.eq_ignore_ascii_case("test") {
+            return Box::new(TestChatClient::new(system_prompt, max_tokens));
+        }
+    }
     match provider.to_lowercase().as_str() {
         "anthropic" => Box::new(AnthropicClient::new(
             model.to_string(),
@@ -92,3 +99,57 @@ pub fn get_chat_client_implementation(
         _ => panic!("Unsupported provider"),
     }
 }
+
+#[cfg(test)]
+struct TestChatClient {
+    system_prompt: String,
+    max_tokens: usize,
+}
+
+#[cfg(test)]
+impl TestChatClient {
+    fn new(system_prompt: String, max_tokens: usize) -> Self {
+        Self {
+            system_prompt,
+            max_tokens,
+        }
+    }
+}
+
+#[cfg(test)]
+impl ChatClient for TestChatClient {
+    fn generate_response(
+        &self,
+        _history_messages_json: Value,
+        _user_prompt: &str,
+        _context_content: Option<&str>,
+    ) -> io::Result<ChatResponse> {
+        Ok(ChatResponse {
+            content: String::from("test-response"),
+            tool_calls: None,
+        })
+    }
+
+    fn generate_tool_response(&self, _tool_prompt: Value) -> io::Result<ChatResponse> {
+        Ok(ChatResponse {
+            content: String::from("test-tool-response"),
+            tool_calls: None,
+        })
+    }
+
+    fn model_context_size(&self) -> Option<usize> {
+        Some(self.max_tokens)
+    }
+
+    fn model_supports_tools(&self) -> bool {
+        false
+    }
+
+    fn update_system_prompt(&mut self, system_prompt: String) {
+        self.system_prompt = system_prompt;
+    }
+
+    fn system_prompt(&self) -> String {
+        self.system_prompt.clone()
+    }
+}
diff --git a/src/main.rs b/src/main.rs
@@ -21,15 +21,15 @@ mod models;
 pub mod tool;
 pub mod traits;
 mod user_input;
+mod session;
 
 #[cfg(test)]
 mod test_support;
 
-use crate::api::{ChatClient, get_chat_client_implementation};
-use crate::command::commands::{CommandResult, create_command_registry};
-use crate::config::{AppConfig, args::Args};
+use crate::command::commands::{create_command_registry, CommandResult};
+use crate::config::{args::Args, AppConfig};
 use crate::models::context_file::ContextFile;
-use crate::models::history_file::HistoryFile;
+use crate::session::Session;
 use crate::traits::estimate_context_size::ContextEstimation;
 use clap::Parser;
 use colored::Colorize;
@@ -38,55 +38,41 @@ use std::io::{self};
 
 fn main() -> io::Result<()> {
     let args = Args::parse();
-    let mut app_config = AppConfig::load_config(&args.history_file);
+    let app_config = AppConfig::load_config(&args.history_file);
     let command_registry = create_command_registry(app_config.user_config.command_prefixes.clone());
     let mut context_file_path = args.context_file.clone();
-    let mut history = HistoryFile::new(
-        app_config.cache_config.get_history_file_path(),
-        app_config.data_dir.display().to_string(),
-    )?;
-
-    history.print_content();
-    app_config.print_model_info();
-
-    let mut chat_client: Box<dyn ChatClient> = get_chat_client_implementation(
-        &app_config.current_profile.provider,
-        &app_config.current_model.model,
-        app_config.user_config.system_prompt.clone(),
-        app_config.user_config.max_tokens,
-    );
+    let mut session = Session::new(app_config);
+
+    session.history_file.print_content();
+    session.config.print_model_info();
+
     let mut rebuild_chat_client = false;
 
     loop {
         if rebuild_chat_client {
-            chat_client = get_chat_client_implementation(
-                &app_config.current_profile.provider,
-                &app_config.current_model.model,
-                app_config.user_config.system_prompt.clone(),
-                app_config.user_config.max_tokens,
-            );
+            session.rebuild_client();
             rebuild_chat_client = false;
         }
 
         // TODO: This shouldn't be printed on every iteration and model information should be fetched once
-        if &app_config.current_profile.provider == "ollama" && chat_client.model_supports_tools() {
+        if session.config.current_profile.provider == "ollama" && session.client.model_supports_tools() {
             println!("Model supports tools");
         }
 
         let context_file = ContextFile::new(&context_file_path);
 
-        if app_config.user_config.token_estimation {
+        if session.config.user_config.token_estimation {
             print_token_usage(
-                history.estimate_context_size() + context_file.estimate_context_size(),
-                &chat_client.model_context_size(),
+                session.history_file.estimate_context_size() + context_file.estimate_context_size(),
+                &session.client.model_context_size(),
             );
         }
 
         println!(
             "\nEnter your prompt or a command (type ':q' to end or ':help' for other command)"
         );
 
-        let mut rl = match app_config.create_rustyline_editor(&command_registry) {
+        let mut rl = match session.config.create_rustyline_editor(&command_registry) {
             Ok(r) => r,
             Err(e) => {
                 eprintln!("Error initializing rustyline: {e}");
@@ -103,9 +89,9 @@ fn main() -> io::Result<()> {
         };
 
         let mut processor = CommandProcessor::new(
-            &mut chat_client,
-            &mut history,
-            &mut app_config,
+            &mut session.client,
+            &mut session.history_file,
+            &mut session.config,
             &command_registry,
             &mut context_file_path,
             &mut rebuild_chat_client,
diff --git a/src/session.rs b/src/session.rs
@@ -0,0 +1,162 @@
+/*
+ * Copyright © 2025 Mitja Leino
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+ * documentation files (the “Software”), to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+ * WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
+ * OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+use crate::api::get_chat_client_implementation;
+use crate::config::AppConfig;
+use crate::models::history_file::HistoryFile;
+
+pub(crate) struct Session {
+    pub(crate) config: AppConfig,
+    pub(crate) client: Box<dyn crate::api::ChatClient>,
+    pub(crate) history_file: HistoryFile,
+}
+
+impl Session {
+    pub(crate) fn new(config: AppConfig) -> Self {
+        let client = Self::create_client(&config);
+        let history_file = HistoryFile::new(
+            config.cache_config.get_history_file_path(),
+            config.data_dir.display().to_string(),
+        ).expect("Failed to determine history file");
+
+        Self {
+            config,
+            client,
+            history_file,
+        }
+    }
+
+    fn create_client(config: &AppConfig) -> Box<dyn crate::api::ChatClient> {
+        get_chat_client_implementation(
+            &config.current_profile.provider,
+            &config.current_model.model,
+            config.user_config.system_prompt.clone(),
+            config.user_config.max_tokens,
+        )
+    }
+
+    pub(crate) fn rebuild_client(&mut self) {
+        self.client = Self::create_client(&self.config);
+    }
+
+    pub fn update_config<F>(&mut self, modifier: F)
+    where
+        F: FnOnce(&mut AppConfig) -> bool,
+    {
+        let should_rebuild = modifier(&mut self.config);
+        if should_rebuild {
+            self.rebuild_client();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Session;
+    use crate::config::profiles_config::{Model, ModelType, Profile};
+    use crate::config::{AppConfig, CacheConfig, UserConfig};
+    use tempfile::TempDir;
+
+    fn make_test_config(temp: &TempDir, history_name: &str, system_prompt: &str) -> AppConfig {
+        let mut config = AppConfig::default();
+
+        // point data_dir to a temporary location to avoid touching user directories
+        let chats_dir = temp.path().join("chats");
+        std::fs::create_dir_all(&chats_dir).unwrap();
+        config.data_dir = chats_dir;
+
+        // minimal cache config so HistoryFile path is resolvable
+        config.cache_config = CacheConfig {
+            last_history_file: Some(history_name.to_string()),
+            last_profile_name: Some("test-profile".to_string()),
+            profile_models: None,
+        };
+
+        // minimal user config
+        let mut user = UserConfig::default();
+        user.system_prompt = system_prompt.to_string();
+        config.user_config = user.clone();
+
+        // set a test profile that maps to our test stub ChatClient
+        config.current_profile = Profile {
+            name: "test-profile".to_string(),
+            provider: "test".to_string(),
+            models: vec![Model {
+                model: "test-model".to_string(),
+                description: Some("test model".to_string()),
+                model_type: ModelType::Fast,
+            }],
+        };
+
+        config.current_model = Model {
+            model: "test-model".to_string(),
+            description: Some("test model".to_string()),
+            model_type: ModelType::Fast,
+        };
+
+        config
+    }
+
+    #[test]
+    fn session_new_initializes_history_file() {
+        let temp = TempDir::new().unwrap();
+        let config = make_test_config(&temp, "sess.hist", "sp-1");
+
+        let session = Session::new(config);
+
+        // History file filename should match
+        assert_eq!(session.history_file.filename, "sess.hist");
+        // File should exist on disk
+        assert!(std::path::Path::new(&session.history_file.path).exists());
+    }
+
+    #[test]
+    fn update_config_without_rebuild_keeps_client() {
+        let temp = TempDir::new().unwrap();
+        let mut session = Session::new(make_test_config(&temp, "sess2.hist", "sp-initial"));
+
+        // Changing the system prompt but returning false should not rebuild client
+        session.update_config(|c| {
+            c.user_config.system_prompt = "sp-updated".to_string();
+            false
+        });
+
+        assert_eq!(session.client.system_prompt(), "sp-initial".to_string());
+    }
+
+    #[test]
+    fn update_config_with_rebuild_updates_client() {
+        let temp = TempDir::new().unwrap();
+        let mut session = Session::new(make_test_config(&temp, "sess3.hist", "sp-a"));
+
+        session.update_config(|c| {
+            c.user_config.system_prompt = "sp-b".to_string();
+            true // request rebuild
+        });
+
+        assert_eq!(session.client.system_prompt(), "sp-b".to_string());
+    }
+
+    #[test]
+    fn rebuild_client_applies_current_config() {
+        let temp = TempDir::new().unwrap();
+        let mut session = Session::new(make_test_config(&temp, "sess4.hist", "sp-0"));
+
+        session.config.user_config.system_prompt = "sp-1".to_string();
+        session.rebuild_client();
+
+        assert_eq!(session.client.system_prompt(), "sp-1".to_string());
+    }
+}
