use sessions instead of cache

src/api.rs

@@ -5,23 +5,29 @@ use std::process::exit;
 use reqwest::{header::{HeaderMap, HeaderValue}, Client};
 use serde::{Deserialize, Serialize};
 
-use crate::{cache::Cache, config::Config};
+use crate::{config::Config, session::Session};
 use crate::{WARN, RED, RESET};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-struct ChatMessagePayload {
+pub struct ChatMessagePayload {
     pub role: String,
     pub content: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-struct ChatPayload {
+pub struct ChatPayload {
     pub model: String,
     pub messages: Vec<ChatMessagePayload>
 }
 
+impl ChatMessagePayload {
+    pub fn is_ai(&self) -> bool {
+        self.role == "assistant"
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
-struct ChatChunk {
+pub struct ChatChunk {
     pub role: Option<String>,
     pub message: String,
     pub created: u64,
@@ -30,8 +36,17 @@ struct ChatChunk {
     pub model: Option<String>
 }
 
+impl Into<ChatMessagePayload> for ChatChunk {
+    fn into(self) -> ChatMessagePayload {
+        ChatMessagePayload {
+            role: "assistant".to_string(),
+            content: self.message
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
-struct ErrChatChunk {
+pub struct ErrChatChunk {
     pub action: String,
     pub status: u64,
     #[serde(rename = "type")]
@@ -87,10 +102,15 @@ pub async fn get_vqd(cli: &Client) -> Result<String, Box<dyn Error>> {
     }
 }
 
-pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, cache: &'a mut Cache, config: &Config) {
+pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, config: &Config) {
+    let init_msg = ChatMessagePayload { role: "user".into(), content: query };
+
+    let mut session = Session::create_or_restore(&vqd);
+    session.push_message(&init_msg).unwrap();
+
     let payload = ChatPayload {
         model: config.model.to_string(),
-        messages: vec![ ChatMessagePayload { role: "user".into(), content: query } ]
+        messages: session.get_messages()
     };
     let payload = serde_json::to_string(&payload).unwrap();
 
@@ -105,15 +125,16 @@ pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, cache: &'a mu
     let new_vqd = res.headers().iter().find(|x| x.0 == "x-vqd-4");
     let vqd_set_res = 
         if let Some(new_vqd) = new_vqd {
-            cache.set_last_vqd(new_vqd.1.as_bytes().iter().map(|x| char::from(*x)).collect::<String>())
+            session.set_last_vqd(new_vqd.1.as_bytes().iter().map(|x| char::from(*x)).collect::<String>())
         } else {
             eprintln!("{WARN}Warn: DuckDuckGo did not return new VQD. Ignore this if everything else is ok.{RESET}");
-            cache.set_last_vqd(vqd.clone())
+            session.set_last_vqd(vqd.clone())
         };
     if let Err(err) = vqd_set_res {
         eprintln!("{WARN}Warn: Could not save VQD to cache: {err}{RESET}");
     }
 
+    let mut error = None;
     while let Some(chunk) = res.chunk().await.unwrap() {
 
         if let Ok(obj) = serde_json::from_slice::<ErrChatChunk>(&chunk) {
@@ -127,9 +148,16 @@ pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, cache: &'a mu
         let chunk = chunk.replace("data: ", "");
         for line in chunk.lines() {
             if let Ok(obj) = serde_json::from_str::<ChatChunk>(line) {
+                if let Err(err) = session.push_ai_message_chunk(&obj) {
+                    error = Some(err);
+                }
                 print!("{}", obj.message);
             }
         }
     }
+    if let Some(err) = error {
+        eprintln!("Error while writing to session: {err:#?}");
+        eprintln!("Session may be broken.");
+    }
     println!("\n");
 }

src/cache.rs

@@ -1,81 +0,0 @@
-
-use std::{env, error::Error, fs, io, path::PathBuf};
-use home::home_dir;
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Cache {
-    pub last_vqd: String,
-    pub last_vqd_time: u64
-}
-
-impl Default for Cache {
-    fn default() -> Self {
-        Self {
-            last_vqd: "".into(),
-            last_vqd_time: 0
-        }
-    }
-}
-
-impl Cache {
-
-    pub fn get_path<T: From<String>>() -> T {
-        match env::var("HEY_CACHE_PATH") {
-            Ok(v) => v,
-            Err(_) =>
-                match home_dir() {
-                    Some(home) => home.join(".cache/hey").as_os_str().as_encoded_bytes().iter().map(|x| char::from(*x)).collect::<String>(),
-                    None => panic!("Cannot detect your home directory!")
-                }
-        }.into()
-    }
-
-    pub fn get_file_name<T: From<String>>() -> T {
-        match env::var("HEY_CACHE_FILENAME") {
-            Ok(v) => v,
-            Err(_) => "cache.json".into()
-        }.into()
-    }
-
-    fn ensure_dir_exists() -> io::Result<()> {
-        let path: PathBuf = Self::get_path();
-        if ! path.is_dir() { fs::create_dir_all(path)? }
-        Ok(())
-    }
-
-    pub fn load() -> Result<Self, Box<dyn Error>> {
-        let path: PathBuf = Self::get_path();
-        Self::ensure_dir_exists()?;
-
-        let file_path = path.join(Self::get_file_name::<PathBuf>());
-        if ! file_path.is_file() {
-            let def = Self::default();
-            def.save()?;
-            Ok(def)
-        } else {
-            let file = fs::read_to_string(file_path)?;
-            Ok(serde_json::from_str(&file)?)
-        }
-    }
-
-    pub fn save(self: &Self) -> Result<(), Box<dyn Error>> {
-        let path: PathBuf = Self::get_path();
-        Self::ensure_dir_exists()?;
-
-        let file_path = path.join(Self::get_file_name::<PathBuf>());
-        fs::write(file_path, serde_json::to_string_pretty(self)?)?;
-        Ok(())
-    }
-
-    pub fn set_last_vqd<T: Into<String>>(self: &mut Self, vqd: T) -> Result<(), Box<dyn Error>> {
-        self.last_vqd = vqd.into();
-        self.last_vqd_time = chrono::Local::now().timestamp_millis() as u64;
-        self.save()?;
-        Ok(())
-    }
-
-    pub fn get_last_vqd<'a, T: From<&'a String>>(self: &'a Self) -> Option<T> {
-        None
-    }
-}

src/config.rs

@@ -1,4 +1,6 @@
-use std::{env, error::Error, fs, io, path::PathBuf};
+use crate::{RED, BLUE, WARN, RESET, GRAY};
+
+use std::{env, error::Error, fs, io, path::PathBuf, process::exit};
 
 use home::home_dir;
 use serde::{Deserialize, Serialize};
@@ -48,6 +50,27 @@ impl Default for Config {
 }
 
 impl Config {
+
+    pub fn set_tos(&mut self, tos: bool) -> Result<(), Box<dyn Error>> {
+        self.tos = tos;
+        self.save()?;
+        Ok(())
+    }
+
+    pub fn check_tos(&self) -> () {
+        if ! self.tos {
+            eprintln!("{RED}You need to agree to DuckDuckgo AI Chat TOS to continue.{RESET}");
+            eprintln!("{GRAY}Visit it on this URL: {RESET}{BLUE}https://duck.ai{RESET}");
+            eprintln!();
+            eprintln!("{GRAY}If it doesnt pop up, open the URL in a private window.{RESET}");
+            eprintln!("{GRAY}Running this command is equivalent to clicking \"I Agree\" in the ToS window:{RESET}");
+            eprintln!("{BLUE}  {} --agree-tos{RESET}", std::env::current_exe().unwrap().display());
+            eprintln!();
+            eprintln!("{WARN}Note: if you want to, modify `tos` parameter in {}{RESET}", Config::get_path::<PathBuf>().join(Config::get_file_name::<String>()).display());
+            exit(3);
+        }
+    }
+
     pub fn get_path<T: From<String>>() -> T {
         match env::var("HEY_CONFIG_PATH") {
             Ok(v) => v,

src/main.rs

@@ -1,22 +1,24 @@
-use std::{path::PathBuf, process::exit};
+use std::process::exit;
 
 use reqwest::Client;
 
 use clap::Parser;
+use session::Session;
 use std::io::{stdout, IsTerminal};
 
-use crate::{cache::Cache, config::Config};
+use crate::config::Config;
 use crate::api::{get_res, get_vqd, simulate_browser_reqs};
 
-mod cache;
 mod config;
 mod api;
+mod session;
 
 pub const GREEN: &str = "\x1b[1;32m";
 pub const RED:   &str = "\x1b[1;31m";
 pub const BLUE:  &str = "\x1b[34m";
 pub const WARN:  &str = "\x1b[33m";
 pub const RESET: &str = "\x1b[0m";
+pub const GRAY:  &str = "\x1b[30m";
 
 #[derive(Debug, Parser)]
 #[command(version, about, long_about = None)]
@@ -38,38 +40,35 @@ async fn main() {
     }
 
     let args = Args::parse();
-    let query = args.query.join(" ");
+    let query = args.query.join(" ").trim().to_string();
+
+    if query.len() == 0 {
+        exit(0);
+    }
 
-    let mut cache = Cache::load().unwrap();
     let mut config = Config::load().unwrap();
 
     if args.agree_tos {
         if ! config.tos {
             println!("{GREEN}TOS accepted{RESET}");
         }
-        config.tos = true;
+        config.set_tos(true).expect("Error saving config");
-        config.save().expect("Error saving config");
-        exit(0);
     }
 
-    if ! config.tos {
+    config.check_tos();
-        eprintln!("{RED}You need to agree to duckduckgo AI chat TOS to continue.{RESET}");
-        eprintln!("{RED}Visit it on this URL: {RESET}{BLUE}https://duckduckgo.com/?q=duckduckgo+ai+chat&ia=chat{RESET}");
-        eprintln!("Once you read it, pass --agree-tos parameter to agree.");
-        eprintln!("{WARN}Note: if you want to, modify `tos` parameter in {}{RESET}", Config::get_path::<PathBuf>().join(Config::get_file_name::<String>()).display());
-        exit(3);
-    }
 
     println!("{GREEN}Contacting DuckDuckGo chat AI...{RESET}");
 
     let cli = Client::new();
     simulate_browser_reqs(&cli).await.unwrap();
 
-    let vqd = match cache.get_last_vqd() {
+    let vqd = match Session::restore_vqd() {
-        Some(v) => { println!("using cached vqd"); v},
+        Some(v) => { v },
         None => get_vqd(&cli).await.unwrap()
     };
 
-    get_res(&cli, query, vqd, &mut cache, &config).await;
+    println!("{vqd:?}");
+
+    get_res(&cli, query, vqd, &config).await;
 
 }

src/session.rs

@@ -0,0 +1,153 @@
+use std::fs;
+use std::path::PathBuf;
+use std::process::{Command, Stdio};
+use std::time::Duration;
+
+use chrono::{DateTime, Local};
+use serde::{Deserialize, Serialize};
+
+use crate::api::{ChatChunk, ChatMessagePayload};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Session {
+    session_id: String,
+    ttl: DateTime<Local>,
+    messages: Vec<ChatMessagePayload>,
+    vqd: String
+}
+
+impl Session {
+    pub fn terminal_session_id<T: From<String>>() -> T {
+        #[cfg(unix)]
+        fn inner() -> Result<String, Box<dyn std::error::Error>> {
+            let child = Command::new("tty")
+                .stdout(Stdio::piped())
+                .stderr(Stdio::null())
+                .spawn()?;
+
+            let child = child.wait_with_output()?;
+            if ! child.status.success() {
+                Err("")? // value ignored
+            }
+
+            Ok(String::from_utf8(child.stdout)?)
+        }
+
+        #[cfg(not(unix))]
+        fn inner() -> Result<String, ()> {
+            Err(())
+        }
+
+        inner()
+            .unwrap_or("default".into())
+                .replace(['/', '\\', ':'], "-")
+                .trim()
+                .trim_start_matches('-')
+                .to_string()
+            .into()
+    }
+
+    fn path_for_id<T: Into<String>>(id: T) -> PathBuf {
+        let id: String = id.into();
+        std::env::temp_dir().join(format!("hey-duck-session-{id}"))
+    }
+    
+    fn path(&self) -> PathBuf {
+        Self::path_for_id(&self.session_id)
+    }
+
+    fn restore_with_id<T: Into<String>>(id: T) -> Option<Self> {
+        fn inner(path: PathBuf) -> Result<Session, Box<dyn std::error::Error>> {
+            if path.is_file() {
+                let data = fs::read_to_string(path)?;
+                Ok(serde_json::from_str(&data)?)
+            } else {
+                Err("")? // value ignored
+            }
+        }
+
+        match inner(Self::path_for_id(id)) {
+            Ok(mut session) => {
+                if session.is_expired() {
+                    session.destroy().expect("Couldn't destroy expired session");
+                    None
+                } else {
+                    session.increase_ttl().expect("Couldn't increase TTL");
+                    Some(session)
+                }
+            },
+            Err(_) => None
+        }
+    }
+
+    fn get_ttl() -> DateTime<Local> {
+        Local::now() + Duration::from_secs(60 * 5)
+    }
+
+    pub fn increase_ttl(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        self.ttl = Self::get_ttl();
+        Ok(self.save()?)
+    }
+
+    fn new<I: Into<String>, V: Into<String>>(id: I, vqd: V) -> Self {
+        Self {
+            session_id: id.into(),
+            ttl: Self::get_ttl(),
+            messages: vec![],
+            vqd: vqd.into()
+        }
+    }
+
+    fn save(&self) -> Result<(), Box<dyn std::error::Error>> {
+        fs::write(self.path(), serde_json::to_string_pretty(self)?)?;
+        Ok(())
+    }
+
+    fn is_expired(&self) -> bool {
+        self.ttl < Local::now()
+    }
+
+    pub fn destroy(&self) -> Result<(), Box<dyn std::error::Error>> {
+        Ok(fs::remove_file(self.path())?)
+    }
+
+    pub fn create_or_restore<T: Into<String>>(vqd: T) -> Self {
+        let session_id: String = Self::terminal_session_id();
+        match Self::restore_with_id(&session_id) {
+            Some(session) => session,
+            None => {
+                let session = Self::new(&session_id, vqd);
+                session.save().expect("Couldn't save new session");
+                session
+            }
+        }
+    }
+
+    pub fn push_message(&mut self, msg: &ChatMessagePayload) -> Result<(), Box<dyn std::error::Error>> {
+        self.messages.push(msg.clone());
+        Ok(self.save()?)
+    }
+
+    pub fn push_ai_message_chunk(&mut self, chunk: &ChatChunk) -> Result<(), Box<dyn std::error::Error>> {
+        if self.messages.last().unwrap().is_ai() {
+            self.messages.last_mut().unwrap().content += chunk.message.as_str();
+        } else {
+            self.messages.push(chunk.clone().into());
+        }
+
+        Ok(self.save()?)
+    }
+
+    pub fn get_messages(&self) -> Vec<ChatMessagePayload> {
+        self.messages.clone()
+    }
+
+    pub fn restore_vqd() -> Option<String> {
+       Self::restore_with_id(Self::terminal_session_id::<String>()).map(|x| x.vqd)
+    }
+
+    pub fn set_last_vqd<T: Into<String>>(&mut self, vqd: T) -> Result<(), Box<dyn std::error::Error>> {
+        self.vqd = vqd.into();
+        Ok(self.save()?)
+    }
+}