From 42deaa082fe026b955f1bc11b5168a2309aad99b Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 5 Nov 2024 17:16:01 +0800
Subject: feat: support session-scoped agent variables (#969)

---
 src/config/session.rs | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

(limited to 'src/config/session.rs')

diff --git a/src/config/session.rs b/src/config/session.rs
index 7335128..8e82c9a 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -27,12 +27,15 @@ pub struct Session {
     #[serde(skip_serializing_if = "Option::is_none")]
     compress_threshold: Option<usize>,
 
+    #[serde(skip_serializing_if = "Option::is_none")]
+    role_name: Option<String>,
+    #[serde(skip_serializing_if = "IndexMap::is_empty")]
+    agent_variables: IndexMap<String, String>,
+
     #[serde(default, skip_serializing_if = "HashMap::is_empty")]
     data_urls: HashMap<String, String>,
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     compressed_messages: Vec<Message>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    role_name: Option<String>,
 
     messages: Vec<Message>,
 
@@ -79,10 +82,6 @@ impl Session {
             }
         }
 
-        if let Some(agent) = &config.agent {
-            session.set_agent(agent);
-        }
-
         Ok(session)
     }
 
@@ -249,16 +248,24 @@ impl Session {
         self.dirty = true;
     }
 
-    pub fn set_agent(&mut self, agent: &Agent) {
-        self.role_prompt
-            .clone_from(&agent.interpolated_instructions());
-    }
-
     pub fn clear_role(&mut self) {
         self.role_name = None;
         self.role_prompt.clear();
     }
 
+    pub fn sync_agent(&mut self, agent: &Agent, set_dirty: bool) {
+        self.role_name = None;
+        self.role_prompt = agent.interpolated_instructions();
+        self.agent_variables = agent.variables().clone();
+        if set_dirty {
+            self.dirty = true;
+        }
+    }
+
+    pub fn agent_variables(&self) -> &IndexMap<String, String> {
+        &self.agent_variables
+    }
+
     pub fn set_save_session(&mut self, value: Option<bool>) {
         if self.name == TEMP_SESSION_NAME {
             return;
-- 
cgit v1.2.3