Agent chat: queue messages while the agent works, Send now, Stop, New chat, and a model picker (#1566, #1571) (#1566)

Thanks @bherbruck.

crates/ui-egui/src/agent_chat/api_client.rs

@@ -14,6 +14,7 @@ use std::time::Duration;
 
 use serde_json::{Value, json};
 
+use super::acp::ModelInfo;
 use super::backend::{ChatBackend, Event, HttpFn, HttpRequest, PromptPart, StartContext, ToolStatus};
 use super::process::{AgentProcess, FromAgent, ProcessHandle};
 
@@ -27,6 +28,8 @@ const MAX_HISTORY: usize = 240;
 const REQUEST_TIMEOUT_S: u64 = 300;
 /// How long one tool call may take (it may wait for the user's Allow).
 const TOOL_TIMEOUT: Duration = Duration::from_secs(330);
+/// How often a waiting request or tool call looks at the cancel flag.
+const CANCEL_POLL: Duration = Duration::from_millis(40);
 
 /// Tools the model is not offered.
 const HIDDEN_TOOLS: &[&str] = &["crash_reports"];
@@ -41,6 +44,8 @@ pub struct ApiConfig {
 
 enum Cmd {
     Prompt(Vec<PromptPart>),
+    NewChat,
+    SetModel(String),
 }
 
 pub struct ApiBackend {
@@ -90,6 +95,18 @@ impl ChatBackend for ApiBackend {
     fn cancel(&mut self) {
         self.cancel.store(true, Ordering::Relaxed);
     }
+    fn set_model(&mut self, model: &str) {
+        if let Some(c) = &self.cmds {
+            let _ = c.send(Cmd::SetModel(model.to_string()));
+        }
+    }
+    fn new_session(&mut self) {
+        // A running turn is cancelled; the worker resets the history once it has ended.
+        self.cancel.store(true, Ordering::Relaxed);
+        if let Some(c) = &self.cmds {
+            let _ = c.send(Cmd::NewChat);
+        }
+    }
     fn answer_permission(&mut self, _key: u64, _option: Option<&str>) {}
     fn authenticate(&mut self, _method: &str) {}
     fn pid(&self) -> Option<u32> {
@@ -134,9 +151,25 @@ impl Worker {
         }
         self.history.push(json!({"role": "system", "content": super::acp::SYSTEM_NOTE}));
         self.emit(Event::Ready { agent: format!("{} · {}", self.cfg.name, self.cfg.model), images: true });
+        // The endpoint's models for the picker; an endpoint that does not list them shows none.
+        if let Ok(list) = model_list(&self.cfg.http, &self.cfg.base_url, self.cfg.key.as_deref()) {
+            let mut available = list;
+            if !self.cfg.model.is_empty() && !available.iter().any(|m| m.id == self.cfg.model) {
+                available.insert(0, ModelInfo { id: self.cfg.model.clone(), name: self.cfg.model.clone(), description: String::new() });
+            }
+            if !available.is_empty() {
+                self.emit(Event::Models { current: self.cfg.model.clone(), available });
+            }
+        }
         for cmd in cmds {
             match cmd {
                 Cmd::Prompt(parts) => self.turn(parts),
+                Cmd::SetModel(m) => self.cfg.model = m,
+                Cmd::NewChat => {
+                    self.history.truncate(1);
+                    self.cancel.store(false, Ordering::Relaxed);
+                    self.emit(Event::Ready { agent: format!("{} · {}", self.cfg.name, self.cfg.model), images: true });
+                }
             }
         }
         if let Some(mut m) = self.mcp.take() {
@@ -152,7 +185,12 @@ impl Worker {
         m.send(&json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}));
         let deadline = std::time::Instant::now() + TOOL_TIMEOUT;
         loop {
-            let left = deadline.saturating_duration_since(std::time::Instant::now());
+            // Wait in short slices so Stop and Send now do not wait for a slow tool: its late
+            // answer is skipped later (its id no longer matches).
+            if self.cancel.load(Ordering::Relaxed) {
+                return Err("cancelled by the user".into());
+            }
+            let left = deadline.saturating_duration_since(std::time::Instant::now()).min(CANCEL_POLL);
             match m.from_agent.recv_timeout(left) {
                 Ok(FromAgent::Line(l)) => {
                     let Ok(v) = serde_json::from_str::<Value>(l.trim()) else { continue };
@@ -165,6 +203,7 @@ impl Worker {
                     return Ok(v.get("result").cloned().unwrap_or(Value::Null));
                 }
                 Ok(FromAgent::Exited(_)) => return Err("the SolidCraft tools stopped".into()),
+                Err(std::sync::mpsc::RecvTimeoutError::Timeout) if std::time::Instant::now() < deadline => {}
                 Err(_) => return Err("the SolidCraft tools did not answer".into()),
             }
         }
@@ -196,12 +235,25 @@ impl Worker {
             self.history.remove(1);
         }
         for _ in 0..MAX_REQUESTS_PER_TURN {
+            // The history as of here is complete (every tool call answered): a cancelled request
+            // goes back to it, so no half-finished round is left to confuse the next message.
+            let sound = self.history.len();
             if self.cancel.load(Ordering::Relaxed) {
                 self.emit(Event::TurnEnded { reason: "cancelled".into() });
                 return;
             }
             let reply = match self.request() {
+                Ok(_) if self.cancel.load(Ordering::Relaxed) => {
+                    self.history.truncate(sound);
+                    self.emit(Event::TurnEnded { reason: "cancelled".into() });
+                    return;
+                }
                 Ok(r) => r,
+                Err(_) if self.cancel.load(Ordering::Relaxed) => {
+                    self.history.truncate(sound);
+                    self.emit(Event::TurnEnded { reason: "cancelled".into() });
+                    return;
+                }
                 Err(e) => {
                     self.emit(Event::Error(e));
                     self.emit(Event::TurnEnded { reason: "error".into() });
@@ -282,7 +334,23 @@ impl Worker {
             body: Some(body.to_string().into_bytes()),
             timeout_s: REQUEST_TIMEOUT_S,
         };
-        let (status, bytes) = (self.cfg.http)(&req)?;
+        // The request runs on a thread of its own so that Stop does not wait for the model: a
+        // cancelled one is abandoned (the host's own timeout ends it) and its answer dropped.
+        let (tx, rx) = channel();
+        let http = self.cfg.http.clone();
+        std::thread::spawn(move || {
+            let _ = tx.send(http(&req));
+        });
+        let (status, bytes) = loop {
+            if self.cancel.load(Ordering::Relaxed) {
+                return Err("cancelled by the user".into());
+            }
+            match rx.recv_timeout(CANCEL_POLL) {
+                Ok(r) => break r?,
+                Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
+                Err(_) => return Err("the request ended without an answer".into()),
+            }
+        };
         let v: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
         if status != 200 {
             return Err(http_error(status, &v));
@@ -413,6 +481,32 @@ fn tool_title(name: &str, args: &Value) -> String {
     }
 }
 
+/// `GET {base_url}/models`: the model ids, in the order the endpoint lists them.
+fn model_ids(http: &HttpFn, base_url: &str, headers: Vec<(String, String)>) -> Result<Vec<String>, String> {
+    let req = HttpRequest { method: "GET", url: format!("{base_url}/models"), headers, body: None, timeout_s: 15 };
+    let (status, bytes) = http(&req)?;
+    let v: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
+    if status != 200 {
+        return Err(http_error(status, &v));
+    }
+    Ok(v.get("data")
+        .and_then(Value::as_array)
+        .map(|a| a.iter().take(500).filter_map(|m| m.get("id").and_then(Value::as_str).map(String::from)).collect())
+        .unwrap_or_default())
+}
+
+/// The endpoint's models for the picker.
+pub fn model_list(http: &HttpFn, base_url: &str, key: Option<&str>) -> Result<Vec<ModelInfo>, String> {
+    let mut headers = Vec::new();
+    if let Some(k) = key {
+        if !super::agents::key_safe_for(base_url) {
+            return Err(crate::i18n::tr("ui.ai_key_needs_https").to_string());
+        }
+        headers.push(("Authorization".to_string(), format!("Bearer {k}")));
+    }
+    Ok(model_ids(http, base_url, headers)?.into_iter().map(|id| ModelInfo { name: id.clone(), id, description: String::new() }).collect())
+}
+
 /// Ask an endpoint for its models (Test connection): how many, and whether `model` is one.
 pub fn test_endpoint(http: &HttpFn, base_url: &str, key: Option<&str>, model: &str) -> Result<String, String> {
     let mut headers = Vec::new();
@@ -422,14 +516,8 @@ pub fn test_endpoint(http: &HttpFn, base_url: &str, key: Option<&str>, model: &s
         }
         headers.push(("Authorization".to_string(), format!("Bearer {k}")));
     }
-    let req = HttpRequest { method: "GET", url: format!("{base_url}/models"), headers, body: None, timeout_s: 15 };
-    let (status, bytes) = http(&req)?;
-    let v: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
-    if status != 200 {
-        return Err(http_error(status, &v));
-    }
-    let ids: Vec<&str> =
-        v.get("data").and_then(Value::as_array).map(|a| a.iter().filter_map(|m| m.get("id").and_then(Value::as_str)).collect()).unwrap_or_default();
+    let ids = model_ids(http, base_url, headers)?;
+    let ids: Vec<&str> = ids.iter().map(String::as_str).collect();
     let found = model.is_empty() || ids.contains(&model);
     Ok(crate::i18n::format(if found { "ui.ai_test_ok" } else { "ui.ai_test_no_model" }, &[("count", &ids.len().to_string()), ("model", model)]))
 }
@@ -469,6 +557,9 @@ mod tests {
         let seen = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
         let log = seen.clone();
         let http: HttpFn = Arc::new(move |r: &HttpRequest| {
+            if r.method == "GET" {
+                return Ok((200, b"{\"data\":[]}".to_vec()));
+            }
             log.lock().unwrap().push(serde_json::from_slice(r.body.as_deref().unwrap_or(b"null")).unwrap());
             let a = answers.lock().unwrap().pop().unwrap_or(json!({}));
             Ok((200, a.to_string().into_bytes()))
@@ -494,6 +585,138 @@ mod tests {
         b.shutdown(true);
     }
 
+    /// Cancel while a model request is in flight ends the turn at once (the request is
+    /// abandoned, its late answer dropped), leaves no half-finished round in the history and
+    /// the next message goes through.
+    #[test]
+    fn cancelling_a_request_ends_the_turn_and_the_next_message_still_works() {
+        let (release_tx, release_rx) = channel::<()>();
+        let release_rx = std::sync::Mutex::new(release_rx);
+        let seen = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
+        let log = seen.clone();
+        let http: HttpFn = Arc::new(move |r: &HttpRequest| {
+            if r.method == "GET" {
+                return Ok((200, b"{\"data\":[]}".to_vec()));
+            }
+            let n = {
+                let mut l = log.lock().unwrap();
+                l.push(serde_json::from_slice(r.body.as_deref().unwrap_or(b"null")).unwrap());
+                l.len()
+            };
+            if n == 1 {
+                // The first request hangs until the test lets it go, then answers late.
+                let _ = release_rx.lock().unwrap().recv_timeout(Duration::from_secs(20));
+                return Ok((200, json!({"choices": [{"message": {"role": "assistant", "content": "Too late."}}]}).to_string().into_bytes()));
+            }
+            Ok((200, json!({"choices": [{"message": {"role": "assistant", "content": "Steered."}}]}).to_string().into_bytes()))
+        });
+        let cfg = ApiConfig { name: "Test".into(), base_url: "http://localhost:1/v1".into(), model: "m".into(), key: None, http };
+        let ctx = StartContext { cwd: std::env::temp_dir(), bridge: None, tools: None, wake: Box::new(|| {}) };
+        let mut b = ApiBackend::start(cfg, ctx).unwrap();
+        let wait = |b: &mut ApiBackend, events: &mut Vec<Event>| {
+            let start = std::time::Instant::now();
+            while !events.iter().any(|e| matches!(e, Event::TurnEnded { .. })) && start.elapsed() < Duration::from_secs(10) {
+                events.extend(b.poll());
+                std::thread::sleep(Duration::from_millis(10));
+            }
+        };
+        b.prompt(&[PromptPart::Text("long job".into())]);
+        let start = std::time::Instant::now();
+        while seen.lock().unwrap().is_empty() && start.elapsed() < Duration::from_secs(10) {
+            std::thread::sleep(Duration::from_millis(10));
+        }
+        b.cancel();
+        let mut events = Vec::new();
+        wait(&mut b, &mut events);
+        assert!(start.elapsed() < Duration::from_secs(8), "the turn ended while the request still hung");
+        assert!(matches!(events.last(), Some(Event::TurnEnded { reason }) if reason == "cancelled"), "{events:?}");
+        assert!(!events.iter().any(|e| matches!(e, Event::Text { .. } | Event::Error(_))), "{events:?}");
+        // The next message is answered, and the late answer of the first never shows up.
+        b.prompt(&[PromptPart::Text("[interrupted] do this instead".into())]);
+        let mut events = Vec::new();
+        wait(&mut b, &mut events);
+        let _ = release_tx.send(());
+        assert!(events.iter().any(|e| matches!(e, Event::Text { text, .. } if text == "Steered.")), "{events:?}");
+        assert!(matches!(events.last(), Some(Event::TurnEnded { reason }) if reason.is_empty()), "{events:?}");
+        let sent = seen.lock().unwrap();
+        let roles: Vec<&str> = sent[1]["messages"].as_array().unwrap().iter().map(|m| m["role"].as_str().unwrap_or("")).collect();
+        assert_eq!(roles, ["system", "user", "user"], "no reply or tool call of the cancelled round is left");
+        b.shutdown(true);
+    }
+
+    /// New chat resets the model's history to the system note: the next request carries nothing
+    /// of the conversation before it.
+    #[test]
+    fn a_new_chat_resets_the_history_to_the_system_note() {
+        let seen = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
+        let log = seen.clone();
+        let http: HttpFn = Arc::new(move |r: &HttpRequest| {
+            if r.method == "GET" {
+                return Ok((200, b"{\"data\":[]}".to_vec()));
+            }
+            log.lock().unwrap().push(serde_json::from_slice(r.body.as_deref().unwrap_or(b"null")).unwrap());
+            Ok((200, json!({"choices": [{"message": {"role": "assistant", "content": "ok"}}]}).to_string().into_bytes()))
+        });
+        let cfg = ApiConfig { name: "Test".into(), base_url: "http://localhost:1/v1".into(), model: "m".into(), key: None, http };
+        let ctx = StartContext { cwd: std::env::temp_dir(), bridge: None, tools: None, wake: Box::new(|| {}) };
+        let mut b = ApiBackend::start(cfg, ctx).unwrap();
+        let run = |b: &mut ApiBackend, f: &dyn Fn(&mut ApiBackend), want: usize| {
+            f(b);
+            let mut events = Vec::new();
+            let start = std::time::Instant::now();
+            while events.iter().filter(|e| matches!(e, Event::TurnEnded { .. } | Event::Ready { .. })).count() < want
+                && start.elapsed() < Duration::from_secs(10)
+            {
+                events.extend(b.poll());
+                std::thread::sleep(Duration::from_millis(10));
+            }
+        };
+        run(&mut b, &|b| b.prompt(&[PromptPart::Text("one".into())]), 2);
+        run(&mut b, &|b| b.new_session(), 1);
+        run(&mut b, &|b| b.prompt(&[PromptPart::Text("two".into())]), 1);
+        let sent = seen.lock().unwrap();
+        let roles: Vec<&str> = sent[1]["messages"].as_array().unwrap().iter().map(|m| m["role"].as_str().unwrap_or("")).collect();
+        assert_eq!(roles, ["system", "user"]);
+        assert_eq!(sent[1]["messages"][1]["content"], "two");
+        b.shutdown(true);
+    }
+
+    /// The endpoint's `/models` fills the picker (with the current model first if it is not
+    /// listed), and choosing one changes the model of the next request.
+    #[test]
+    fn the_endpoint_lists_models_and_a_choice_changes_the_next_request() {
+        let seen = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
+        let log = seen.clone();
+        let http: HttpFn = Arc::new(move |r: &HttpRequest| {
+            if r.method == "GET" {
+                assert!(r.url.ends_with("/models"));
+                return Ok((200, json!({"data": [{"id": "big"}, {"id": "small"}]}).to_string().into_bytes()));
+            }
+            log.lock().unwrap().push(serde_json::from_slice(r.body.as_deref().unwrap_or(b"null")).unwrap());
+            Ok((200, json!({"choices": [{"message": {"role": "assistant", "content": "ok"}}]}).to_string().into_bytes()))
+        });
+        let cfg = ApiConfig { name: "Test".into(), base_url: "http://localhost:1/v1".into(), model: "local".into(), key: None, http };
+        let ctx = StartContext { cwd: std::env::temp_dir(), bridge: None, tools: None, wake: Box::new(|| {}) };
+        let mut b = ApiBackend::start(cfg, ctx).unwrap();
+        let mut events = Vec::new();
+        let start = std::time::Instant::now();
+        while !events.iter().any(|e| matches!(e, Event::Models { .. })) && start.elapsed() < Duration::from_secs(10) {
+            events.extend(b.poll());
+            std::thread::sleep(Duration::from_millis(10));
+        }
+        let Some(Event::Models { current, available }) = events.iter().find(|e| matches!(e, Event::Models { .. })) else { panic!("{events:?}") };
+        assert_eq!(current, "local");
+        assert_eq!(available.iter().map(|m| m.id.as_str()).collect::<Vec<_>>(), ["local", "big", "small"]);
+        b.set_model("big");
+        b.prompt(&[PromptPart::Text("hi".into())]);
+        let start = std::time::Instant::now();
+        while seen.lock().unwrap().is_empty() && start.elapsed() < Duration::from_secs(10) {
+            std::thread::sleep(Duration::from_millis(10));
+        }
+        assert_eq!(seen.lock().unwrap()[0]["model"], "big");
+        b.shutdown(true);
+    }
+
     #[test]
     fn tool_arguments_are_an_object_or_an_error() {
         assert_eq!(tool_arguments(None), Ok(json!({})));
@@ -515,6 +738,9 @@ mod tests {
         let seen = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
         let log = seen.clone();
         let http: HttpFn = Arc::new(move |r: &HttpRequest| {
+            if r.method == "GET" {
+                return Ok((200, b"{\"data\":[]}".to_vec()));
+            }
             log.lock().unwrap().push(serde_json::from_slice(r.body.as_deref().unwrap_or(b"null")).unwrap());
             Ok((200, answers.lock().unwrap().pop().unwrap_or(json!({})).to_string().into_bytes()))
         });