diff --git a/apps/maple-agent/README.md b/apps/maple-agent/README.md index a5377d822..d16e848cc 100644 --- a/apps/maple-agent/README.md +++ b/apps/maple-agent/README.md @@ -70,6 +70,11 @@ agent runtime; its desktop Agent page points users to Maple Work. attachments (picker, paste, or drag and drop), a per-message Copy button, and a context-window indicator. Every tool call runs without asking for approval. +- A task can switch models between messages: its model menu offers every + model, a pick applies from the next message, and the transcript notes the + switch. When the new model's context window is smaller, the model that + wrote the history summarizes it before the next message if it no longer + fits, as `crates/maple-agent/src/agent/switch_compaction.rs` describes. - Shell tool calls start in the task's saved folder on the execution computer. A `cd` inside a command does not affect later calls. Repository searches should name the intended checkout/worktree with explicit absolute paths. diff --git a/apps/maple-agent/app/src/ui/chat/composer.rs b/apps/maple-agent/app/src/ui/chat/composer.rs index 837364f1b..3192b6f63 100644 --- a/apps/maple-agent/app/src/ui/chat/composer.rs +++ b/apps/maple-agent/app/src/ui/chat/composer.rs @@ -22,44 +22,6 @@ use crate::ui::theme; use crate::ui::titlebar; use crate::ui::widgets; -/// Research's compact label precedence, projected once when the catalog loads. -/// IDs remain the keys used for selection and persisted tasks. -pub(super) fn catalog_model_labels( - catalog: &maple_sdk::ModelCatalogResponse, -) -> std::collections::HashMap { - let label = |short: Option<&str>, full: Option<&str>, id: &str| { - short - .filter(|value| !value.trim().is_empty()) - .or_else(|| full.filter(|value| !value.trim().is_empty())) - .unwrap_or(id) - .to_string() - }; - catalog - .data - .iter() - .map(|model| { - ( - model.id.clone(), - label( - model.short_name.as_deref(), - model.display_name.as_deref(), - &model.id, - ), - ) - }) - .chain(catalog.aliases.iter().map(|alias| { - ( - alias.id.clone(), - label( - alias.short_name.as_deref(), - alias.label.as_deref(), - &alias.id, - ), - ) - })) - .collect() -} - /// The composer's button row, left to right. The row renders from this /// list and application Vim walks it, so a new control is reachable from /// the keyboard once it has an entry here, a rendering, and an action. @@ -370,7 +332,7 @@ impl ChatScreen { ChatPopup::Model, chip( "model-picker", - self.model_locked.then_some("lock"), + None, label, true, self.popup.is_open(&ChatPopup::Model), @@ -562,27 +524,21 @@ impl ChatScreen { fn model_menu(&self) -> Menu { let selected = self.selected_model.as_deref(); - let locked = self.model_locked; + // A task offers every model; a pick applies from its next message. Menu::new("model-menu", px(320.)) .label("Model") .max_height(px(320.)) .application_vim(self.application_vim_enabled) - .items(self.model_menu_models().into_iter().map(|model| { + .items(self.models.iter().map(|model| { let pick = model.clone(); - let item = MenuItem::new( + MenuItem::new( SharedString::from(format!("model-{model}")), - self.model_label(&model).to_string(), + self.model_label(model).to_string(), move |this: &mut Self, _: &mut Window, cx: &mut Context| { this.pick_model(pick.clone(), cx); }, ) - .current(selected == Some(model.as_str())); - // A task keeps the model of its first message. - if locked { - item.note("Locked to this task. Start a new task to switch.") - } else { - item - } + .current(selected == Some(model.as_str())) })) } @@ -1559,35 +1515,3 @@ fn chip( .child(div().whitespace_nowrap().child(label)) .when(chevron, |el| el.child(icon("chevron-down", px(14.), color))) } - -#[cfg(test)] -mod tests { - use super::catalog_model_labels; - - #[test] - fn model_labels_follow_research_catalog_precedence() { - let catalog = serde_json::from_value(serde_json::json!({ - "object": "list", - "data": [ - {"id": "model-a", "short_name": "Compact", "display_name": "Full name"}, - {"id": "model-b", "display_name": "Server supplied name"}, - {"id": "model-c", "short_name": " ", "display_name": "Fallback name"}, - {"id": "model-d"} - ], - "aliases": [ - {"id": "alias-a", "short_name": "Alias", "label": "Full alias", "target_model": "model-a"}, - {"id": "alias-b", "label": "Automatic", "target_model": "model-b"}, - {"id": "alias-c", "short_name": "", "label": " ", "target_model": "model-c"} - ] - })) - .unwrap(); - let labels = catalog_model_labels(&catalog); - assert_eq!(labels["model-a"], "Compact"); - assert_eq!(labels["model-b"], "Server supplied name"); - assert_eq!(labels["model-c"], "Fallback name"); - assert_eq!(labels["model-d"], "model-d"); - assert_eq!(labels["alias-a"], "Alias"); - assert_eq!(labels["alias-b"], "Automatic"); - assert_eq!(labels["alias-c"], "alias-c"); - } -} diff --git a/apps/maple-agent/app/src/ui/chat/mod.rs b/apps/maple-agent/app/src/ui/chat/mod.rs index 6ca64aee1..f0f51883f 100644 --- a/apps/maple-agent/app/src/ui/chat/mod.rs +++ b/apps/maple-agent/app/src/ui/chat/mod.rs @@ -549,8 +549,6 @@ pub struct ChatScreen { thinking_loading: bool, thinking_generation: u64, thinking_saving: bool, - /// The task on screen has a persisted model after its first message. - model_locked: bool, /// Desktop notifications enabled (settings). notify_enabled: bool, /// Local-store policy, independent of the selected execution computer. @@ -1410,7 +1408,6 @@ impl ChatScreen { thinking_loading: false, thinking_generation: 0, thinking_saving: false, - model_locked: false, notify_enabled: settings.desktop_notifications, notify_details: settings.notification_details_enabled(), window_active: true, @@ -2232,7 +2229,7 @@ impl ChatScreen { this.models_loading = false; if let Ok((mut models, saved, catalog)) = result { if let Some(catalog) = &catalog { - this.model_labels = composer::catalog_model_labels(catalog); + this.model_labels = maple_agent::agent::catalog_model_labels(catalog); // An older execution computer lists every API model, // or every catalog model; offer only what the catalog // describes and does not keep out of Maple Work. @@ -2254,7 +2251,7 @@ impl ChatScreen { .or(saved.filter(offered)) .or_else(|| models.first().cloned()); } - if !this.model_locked && this.selected_model.is_none() { + if this.selected_model.is_none() { this.selected_model = this.default_model.clone(); } this.models = models; @@ -3303,7 +3300,7 @@ impl ChatScreen { self.draft = false; self.draft_mcp_defaults = None; self.draft_mcp_changes.clear(); - // The lock belongs to the task; the empty screen uses the default. + // The empty screen uses the default, not the last task's model. self.restore_default_model(cx); self.sync_sidebar_selection(cx); self.set_queue(Vec::new()); @@ -5510,22 +5507,9 @@ impl ChatScreen { ); } + /// Pick the model of the next message, in a task or a new one. It also + /// becomes the account default for new tasks. fn pick_model(&mut self, model: String, cx: &mut Context) { - if self.model_locked { - self.popup.close(cx); - if self.selected_model.as_deref() != Some(model.as_str()) - && let Some(locked) = self.selected_model.as_deref() - { - self.notice = Some( - format!( - "This task is locked to model {locked}. Start a new task to use {model}." - ) - .into(), - ); - } - cx.notify(); - return; - } self.selected_model = Some(model.clone()); self.default_model = Some(model.clone()); self.refresh_thinking(cx); @@ -5553,22 +5537,21 @@ impl ChatScreen { self.refresh_context_usage(cx); } - /// New Task and unlocked sessions use the account default, not the - /// last opened task's locked model. + /// New Task and drafts use the account default, not the last opened + /// task's model. fn restore_default_model(&mut self, cx: &mut Context) { - self.model_locked = false; self.set_composer_model(self.default_model.clone(), cx); self.refresh_thinking(cx); } + /// A started task opens on the model it last used. fn apply_session_model(&mut self, session: &AgentSessionSummary, cx: &mut Context) { - let locked = session.message_count > 0 && session.model.is_some(); - if locked { - self.model_locked = true; - self.set_composer_model(session.model.clone(), cx); - self.refresh_thinking(cx); - } else { - self.restore_default_model(cx); + match session.model.clone().filter(|_| session.message_count > 0) { + Some(model) => { + self.set_composer_model(Some(model), cx); + self.refresh_thinking(cx); + } + None => self.restore_default_model(cx), } } @@ -5576,14 +5559,6 @@ impl ChatScreen { self.model_labels.get(model).map_or(model, String::as_str) } - fn model_menu_models(&self) -> Vec { - if self.model_locked { - self.selected_model.iter().cloned().collect() - } else { - self.models.clone() - } - } - /// Re-measure one row after it changed in place. The list keeps the /// row's last height as a hint and the scroll position as it was, so /// neither the scroll range nor the view moves before the next paint. @@ -5750,15 +5725,18 @@ impl ChatScreen { if session_summary_eq(existing, &session) { return false; } - let mut selected = None; + let mut switched = None; if self.selected_session.as_deref() == Some(session.id.as_str()) { self.web_enabled = session.web_enabled; - selected = Some(session.clone()); + // Follow a switch the task made, here or on another surface; + // a pick waiting for the next message stays. + if existing.model != session.model { + switched = Some(session.clone()); + } } *existing = session; - // The first message locks the task on screen to its model. - if let Some(selected) = selected { - self.apply_session_model(&selected, cx); + if let Some(switched) = switched { + self.apply_session_model(&switched, cx); } } else { self.sessions.insert(0, session); diff --git a/apps/maple-agent/app/src/ui/chat/tests.rs b/apps/maple-agent/app/src/ui/chat/tests.rs index 210a634fa..343b6dba9 100644 --- a/apps/maple-agent/app/src/ui/chat/tests.rs +++ b/apps/maple-agent/app/src/ui/chat/tests.rs @@ -3198,10 +3198,11 @@ mod state_tests { }); } - /// Opening a task with a persisted model points the selector at that - /// model without changing the account default. New Task restores it. + /// Opening a task points the selector at the model it last used without + /// changing the account default. Its menu offers every model, and a pick + /// becomes the account default, as on New Task. #[gpui::test] - fn test_existing_task_adopts_its_locked_model(cx: &mut TestAppContext) { + fn test_existing_task_opens_on_its_model_and_offers_every_model(cx: &mut TestAppContext) { cx.executor().allow_parking(); let screen = screen(cx); screen.update(cx, |this, cx| { @@ -3221,45 +3222,33 @@ mod state_tests { kimi.model = Some("kimi-k3".into()); kimi.message_count = 3; - this.set_active_session(kimi, Vec::new(), HashMap::new(), cx); + this.set_active_session(kimi.clone(), Vec::new(), HashMap::new(), cx); assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); - assert!(this.model_locked); assert_eq!(this.default_model.as_deref(), Some("glm-5-3")); - assert_eq!(this.model_menu_models(), vec!["kimi-k3".to_string()]); this.pick_model("glm-5-3".into(), cx); - assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); + assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); assert_eq!(this.default_model.as_deref(), Some("glm-5-3")); - assert!( - this.notice - .as_ref() - .is_some_and(|notice| notice.as_ref().contains("locked to model kimi-k3")), - "picking another model on a locked task must not change the selector" - ); + assert!(this.notice.is_none()); + // Reopened before a message, the task is still on its own model. this.set_active_session(glm, Vec::new(), HashMap::new(), cx); - assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); - assert!(this.model_locked); + this.set_active_session(kimi, Vec::new(), HashMap::new(), cx); + assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); this.new_session(cx); assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); - assert!(!this.model_locked); - assert_eq!( - this.model_menu_models(), - vec!["glm-5-3".to_string(), "kimi-k3".to_string()] - ); - this.pick_model("kimi-k3".into(), cx); assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); assert_eq!(this.default_model.as_deref(), Some("kimi-k3")); - assert!(!this.model_locked); }); } - /// A task opened before its first message locks as soon as the runtime - /// reports that message, without being opened again. + /// A pick waits for the task's next message: updates of the task on its + /// old model leave it alone, and a switch the task makes, here or on + /// another surface, moves the selector to the new model. #[gpui::test] - fn test_first_message_locks_the_task_on_screen(cx: &mut TestAppContext) { + fn test_a_pick_waits_for_the_next_message_and_switches_follow(cx: &mut TestAppContext) { cx.executor().allow_parking(); use maple_agent::agent::AgentRunEvent; let screen = screen(cx); @@ -3274,29 +3263,50 @@ mod state_tests { created.message_count = 0; this.upsert_session(created.clone(), cx); this.set_active_session(created.clone(), Vec::new(), HashMap::new(), cx); - assert!(!this.model_locked); let mut sent = created; - sent.message_count = 1; - this.handle_run_event("new-task", "r1", AgentRunEvent::SessionUpdated(sent), cx); - assert!(this.model_locked); - assert_eq!(this.model_menu_models(), vec!["glm-5-3".to_string()]); + sent.message_count = 2; + this.handle_run_event( + "new-task", + "r1", + AgentRunEvent::SessionUpdated(sent.clone()), + cx, + ); + assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); this.pick_model("kimi-k3".into(), cx); - assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); - assert_eq!(this.default_model.as_deref(), Some("glm-5-3")); - assert!( - this.notice - .as_ref() - .is_some_and(|notice| notice.as_ref().contains("locked to model glm-5-3")) + sent.updated_ms += 1; + this.handle_run_event( + "new-task", + "r1", + AgentRunEvent::SessionUpdated(sent.clone()), + cx, ); + assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); + + let mut switched = sent.clone(); + switched.model = Some("kimi-k3".into()); + this.handle_run_event( + "new-task", + "r2", + AgentRunEvent::SessionUpdated(switched.clone()), + cx, + ); + assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); + + // Another surface runs the task on GLM again. + let mut elsewhere = switched; + elsewhere.model = Some("glm-5-3".into()); + this.upsert_session(elsewhere, cx); + assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); + assert_eq!(this.default_model.as_deref(), Some("kimi-k3")); }); } /// Deleting the last task of a project leaves an empty screen that uses - /// the default model, not the deleted task's lock. + /// the default model, not the deleted task's model. #[gpui::test] - fn test_leaving_a_locked_task_restores_the_default_model(cx: &mut TestAppContext) { + fn test_leaving_a_task_restores_the_default_model(cx: &mut TestAppContext) { cx.executor().allow_parking(); let screen = screen(cx); screen.update(cx, |this, cx| { @@ -3310,13 +3320,11 @@ mod state_tests { kimi.message_count = 3; this.sessions = vec![kimi.clone()]; this.set_active_session(kimi, Vec::new(), HashMap::new(), cx); - assert!(this.model_locked); + assert_eq!(this.selected_model.as_deref(), Some("kimi-k3")); this.remove_session("kimi-task", cx); assert!(this.selected_session.is_none()); - assert!(!this.model_locked); assert_eq!(this.selected_model.as_deref(), Some("glm-5-3")); - assert_eq!(this.model_menu_models().len(), 2); let request = this.new_session_request().expect("request"); assert_eq!(request.model.as_deref(), Some("glm-5-3")); }); diff --git a/apps/maple-agent/crates/UPSTREAM.md b/apps/maple-agent/crates/UPSTREAM.md index e6552e42f..f83bda589 100644 --- a/apps/maple-agent/crates/UPSTREAM.md +++ b/apps/maple-agent/crates/UPSTREAM.md @@ -14,6 +14,15 @@ review's PR records its adopt/skip decisions and any new differences. This file records the storage follow-up's adaptations; it does not claim a complete audit of every earlier difference. +## Fidelity to upstream + +The crates mirror upstream Pi's behavior, including its bugs. Change them +locally only for a blocker, and say in the PR why no hook or adapter could +carry the fix. A review finding that matches upstream is recorded in the +review's PR and left as it is. Maple-specific behavior belongs in +`maple-agent`, usually as a Pi extension: `agent/switch_compaction.rs`, for +one, changes which model summarizes through `before_compact`. + ## Storage adaptations | Interface | Reason and boundary | Regression evidence | diff --git a/apps/maple-agent/crates/maple-agent/src/acp/convert.rs b/apps/maple-agent/crates/maple-agent/src/acp/convert.rs index be9c2d20e..912c10a22 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/convert.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/convert.rs @@ -56,19 +56,6 @@ pub(super) fn acp_config_options( ] } -pub(super) fn acp_session_config_options( - model: &str, - available_models: &[String], - message_count: usize, -) -> Vec { - if message_count == 0 { - acp_config_options(model, available_models) - } else { - let locked_models = [model.to_string()]; - acp_config_options(model, &locked_models) - } -} - pub(super) fn add_thinking_option( options: &mut Vec, state: &crate::agent::AgentThinkingState, diff --git a/apps/maple-agent/crates/maple-agent/src/acp/mod.rs b/apps/maple-agent/crates/maple-agent/src/acp/mod.rs index 06cf1ae46..d70cd6b0a 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/mod.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/mod.rs @@ -301,7 +301,6 @@ impl AcpConnectionContext { thinking: None, projection: AcpProjection::default(), advertised_title: Some(created.detail.session.title.clone()), - message_count: created.detail.session.message_count, created_here: true, prompted: false, project_root: project_root.clone(), @@ -593,13 +592,6 @@ impl AcpConnectionContext { .data("The requested task does not exist in the supplied project directory")); } let available_models = self.available_models().await?; - if let Some(model) = persisted.model.as_ref() - && !available_models.iter().any(|available| available == model) - { - return Err(agent_client_protocol::Error::invalid_request().data(format!( - "This Maple Work task uses model '{model}', which is no longer available; the task remains available in Maple Desktop" - ))); - } let session_servers = prepare_session_mcp(&request.mcp_servers)?; let protocol_session_id = SessionId::new(session_id.clone()); // Rebuilding announces HistoryReplaced. Finish before registering this @@ -630,11 +622,14 @@ impl AcpConnectionContext { return Err(agent_client_protocol::Error::invalid_request() .data("The task's project changed while loading")); } + // A task whose model is no longer offered opens on the default and + // switches with its next prompt. let Some(model) = attached .detail .session .model .clone() + .filter(|model| available_models.contains(model)) .or_else(|| available_models.first().cloned()) else { lease.release().await; @@ -642,7 +637,6 @@ impl AcpConnectionContext { .await; return Err(internal_acp_error("Maple returned no models".to_string())); }; - let message_count = attached.detail.session.message_count; let finalization = self.finalization.lock().await; if self.closed.load(Ordering::SeqCst) || operation.cancellation.is_cancelled() { drop(finalization); @@ -679,7 +673,6 @@ impl AcpConnectionContext { thinking: None, projection: AcpProjection::default(), advertised_title: Some(persisted.title.clone()), - message_count, created_here: false, prompted: false, project_root: project_root.clone(), @@ -962,8 +955,6 @@ impl AcpConnectionContext { let session = sessions.get_mut(session_id).ok_or_else(|| { agent_client_protocol::Error::resource_not_found(Some(session_id.into())) })?; - session.message_count = usize::try_from(summary.message_count) - .map_err(|_| paged::replay_error("Invalid Maple task message count"))?; if let Some(model) = summary.model && session.model != model { @@ -1017,10 +1008,7 @@ impl AcpConnectionContext { agent_client_protocol::Error::invalid_params().data("Unknown Maple model") ); } - if session.message_count > 0 && session.model != model { - return Err(agent_client_protocol::Error::invalid_params() - .data("Maple tasks are model-locked after their first message")); - } + // A started task switches with its next prompt. let model = model.to_string(); drop(sessions); self.agent @@ -1628,26 +1616,14 @@ impl AcpConnectionContext { .lock() .await .insert(session_id.to_string(), run.run_id.clone()); - let locked_config_options = { - let mut sessions = self.sessions.lock().await; - sessions.get_mut(session_id).and_then(|session| { - session.prompted = true; - let first_message = session.message_count == 0; - session.message_count = session.message_count.saturating_add(1); - first_message.then(|| session.config_options()) - }) - }; + if let Some(session) = self.sessions.lock().await.get_mut(session_id) { + session.prompted = true; + } self.stats.active_runs.fetch_add(1, Ordering::SeqCst); let mut terminal = run.terminal.clone(); let usage = run.usage.clone(); let result = self - .stream_run( - cx, - &protocol_session_id, - run, - locked_config_options, - cancellation, - ) + .stream_run(cx, &protocol_session_id, run, cancellation) .await; let (result, stopped_early) = match result { Ok((response, stopped)) => (Ok(response), stopped), @@ -1696,34 +1672,12 @@ impl AcpConnectionContext { cx: &ConnectionTo, protocol_session_id: &SessionId, run: crate::agent::AgentRunHandle, - locked_config_options: Option>, cancellation: &CancellationToken, ) -> Result<(PromptResponse, bool), agent_client_protocol::Error> { let session_id = protocol_session_id.0.to_string(); let mut events = run.events; let mut terminal = run.terminal; let overflowed = run.event_overflowed; - if let Some(config_options) = locked_config_options { - let update = SessionUpdate::ConfigOptionUpdate(ConfigOptionUpdate::new(config_options)); - match self - .send_session_update( - cx, - SessionNotification::new(protocol_session_id.clone(), update), - cancellation, - ) - .await - { - Ok(()) => {} - Err(AcpOutboundSendError::Cancelled) => { - return Ok((PromptResponse::new(StopReason::Cancelled), true)); - } - Err(AcpOutboundSendError::UpdateTooLarge) => { - return Err(agent_client_protocol::Error::internal_error() - .data("Maple's locked model selector exceeded the ACP update limit")); - } - Err(AcpOutboundSendError::Transport(error)) => return Err(error), - } - } let operation = self .session_operations .lock() diff --git a/apps/maple-agent/crates/maple-agent/src/acp/observer.rs b/apps/maple-agent/crates/maple-agent/src/acp/observer.rs index 92bb988fb..14d501a71 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/observer.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/observer.rs @@ -448,20 +448,15 @@ impl AcpConnectionContext { let changed = { let mut sessions = self.sessions.lock().await; sessions.get_mut(session_id).and_then(|session| { - let changed = session.message_count != summary.message_count - || summary - .model - .as_ref() - .is_some_and(|model| model != &session.model); - session.message_count = summary.message_count; - if let Some(model) = &summary.model - && model != &session.model - { - session.model = model.clone(); - session.catalog = None; - session.thinking = None; - } - changed.then(|| session.config_options()) + // Follow a model the task switched to elsewhere. + let model = summary + .model + .as_ref() + .filter(|model| **model != session.model)?; + session.model = model.clone(); + session.catalog = None; + session.thinking = None; + Some(session.config_options()) }) }; if let Some(options) = changed { diff --git a/apps/maple-agent/crates/maple-agent/src/acp/session.rs b/apps/maple-agent/crates/maple-agent/src/acp/session.rs index 6c54ed23a..68938748b 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/session.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/session.rs @@ -3,7 +3,7 @@ use super::SharedRuntimeStart; use super::config::{AgentAcpConfig, AgentAcpStats}; -use super::convert::acp_session_config_options; +use super::convert::acp_config_options; use super::observer::ObserverMailbox; use super::transport::AcpOutboundTracker; #[cfg(test)] @@ -67,7 +67,6 @@ pub(super) struct AcpSession { /// The session title last advertised to the caller; semantic-title /// updates are only sent when the title moves past this. pub(super) advertised_title: Option, - pub(super) message_count: usize, pub(super) created_here: bool, pub(super) prompted: bool, pub(super) project_root: PathBuf, @@ -141,8 +140,7 @@ impl AcpSessionOperation { impl AcpSession { pub(super) fn config_options(&self) -> Vec { - let mut options = - acp_session_config_options(&self.model, &self.available_models, self.message_count); + let mut options = acp_config_options(&self.model, &self.available_models); if let Some(state) = &self.thinking { super::convert::add_thinking_option(&mut options, state); } diff --git a/apps/maple-agent/crates/maple-agent/src/acp/tests.rs b/apps/maple-agent/crates/maple-agent/src/acp/tests.rs index d47455f74..ad6154e40 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/tests.rs @@ -3,8 +3,7 @@ use super::config::MAX_ACP_CONNECTIONS; use super::convert::{ - acp_session_config_options, acp_tool_update, client_supports_form_elicitation, - timeline_tool_text, + acp_config_options, acp_tool_update, client_supports_form_elicitation, timeline_tool_text, }; use super::session::completed_runtime_start; use super::transport::is_session_update_line; @@ -816,8 +815,8 @@ async fn set_mode_is_rejected() { } /// A turn streams to the caller as it happens: thinking and text once each, -/// tool calls with their results, the locked model selector and the task's -/// title; it ends with the turn's usage. The task is listed from then on, +/// tool calls with their results and the task's title; it ends with the +/// turn's usage. The task is listed from then on, /// and a later connection replays it. #[tokio::test] async fn a_prompt_streams_its_turn_and_a_later_load_replays_it() { @@ -867,14 +866,8 @@ async fn a_prompt_streams_its_turn_and_a_later_load_replays_it() { .contains("Remember the milk."), "{last}" ); - let locked = client.take("config_option_update"); - assert_eq!( - locked[0]["update"]["configOptions"][0]["options"] - .as_array() - .unwrap() - .len(), - 1 - ); + // The first message changes no option: the model selector stays. + assert!(client.take("config_option_update").is_empty()); let titles = client.take("session_info_update"); assert_eq!(titles[0]["update"]["title"], "What do the notes say?"); // Prompted, the task is listed. @@ -905,6 +898,55 @@ async fn a_prompt_streams_its_turn_and_a_later_load_replays_it() { finish(serving).await; } +/// A started task takes another model: the selector offers every model +/// after the first prompt, and the next prompt runs on the one picked. +#[tokio::test] +async fn a_started_task_switches_to_the_model_picked() { + let harness = Harness::new().await; + harness + .serve_catalog(json!({ + "object": "list", + "data": [{"id": "glm-5-3"}, {"id": "kimi-k3"}], + "aliases": [], + })) + .await; + harness.faux.push_text("First."); + harness.faux.push_text("Second."); + let (mut client, serving) = connect(&harness.handle, no_answers()); + client.initialize(json!({})).await; + let session_id = client.new_session(harness.project.path()).await; + let first = client.prompt(&session_id, "One").await; + assert_eq!(first["stopReason"], "end_turn"); + + let picked = client + .call( + "session/set_config_option", + json!({"sessionId": session_id, "configId": "model", "value": "kimi-k3"}), + ) + .await; + let selector = &picked["configOptions"][0]; + assert_eq!(selector["currentValue"], "kimi-k3", "{picked}"); + assert_eq!(selector["options"].as_array().unwrap().len(), 2); + let second = client.prompt(&session_id, "Two").await; + assert_eq!(second["stopReason"], "end_turn"); + + let models: Vec = harness + .faux + .requests() + .iter() + .map(|request| request.model.id.clone()) + .collect(); + assert_eq!(models, ["glm-5-3", "kimi-k3"]); + let summary = harness + .handle + .task_summary(session_id.clone()) + .await + .unwrap(); + assert_eq!(summary.model.as_deref(), Some("kimi-k3")); + client.shutdown().await; + finish(serving).await; +} + /// `session/cancel` stops a turn in progress; the next prompt runs. #[tokio::test] async fn cancel_stops_the_turn_in_progress() { @@ -1502,26 +1544,20 @@ fn prompt_usage_serializes_one_turn_without_session_accumulation() { assert_eq!(encoded["cachedWriteTokens"], 1); } -/// The model selector is the only config option, and it locks to the -/// persisted model after the first message. +/// The model selector is the only config option, and offers every model. #[test] -fn model_selector_locks_to_the_persisted_model_after_first_message() { +fn model_selector_offers_every_model() { let models = vec!["model-a".to_string(), "model-b".to_string()]; - let fresh = serde_json::to_value(acp_session_config_options("model-b", &models, 0)).unwrap(); - let locked = serde_json::to_value(acp_session_config_options("model-b", &models, 1)).unwrap(); + let options = serde_json::to_value(acp_config_options("model-b", &models)).unwrap(); assert_eq!( - fresh.as_array().unwrap().len(), + options.as_array().unwrap().len(), 1, - "no mode option: {fresh}" - ); - assert_eq!(fresh[0]["id"], "model"); - assert_eq!(fresh[0]["currentValue"], "model-b"); - assert_eq!(fresh[0]["options"].as_array().unwrap().len(), 2); - assert_eq!(locked.as_array().unwrap().len(), 1); - assert_eq!(locked[0]["currentValue"], "model-b"); - assert_eq!(locked[0]["options"].as_array().unwrap().len(), 1); - assert_eq!(locked[0]["options"][0]["value"], "model-b"); + "no mode option: {options}" + ); + assert_eq!(options[0]["id"], "model"); + assert_eq!(options[0]["currentValue"], "model-b"); + assert_eq!(options[0]["options"].as_array().unwrap().len(), 2); } #[test] diff --git a/apps/maple-agent/crates/maple-agent/src/agent.rs b/apps/maple-agent/crates/maple-agent/src/agent.rs index 4b6144975..9d37a8024 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent.rs @@ -30,6 +30,7 @@ mod runtime; mod side_models; mod store; mod surface; +mod switch_compaction; mod tasks; mod thinking; pub(crate) mod timeline; @@ -43,6 +44,7 @@ use std::sync::{Arc, RwLock}; use crate::maple_api::{MapleApiSession, account_scope}; pub use attachments::{AgentAttachment, AgentImageUpload}; +pub use catalog::catalog_model_labels; use config::*; pub use config::{ MAPLE_WORKSPACE_DIRECTORY_NAME, MAPLE_WORKSPACE_DISPLAY_NAME, default_maple_workspace_path, diff --git a/apps/maple-agent/crates/maple-agent/src/agent/catalog.rs b/apps/maple-agent/crates/maple-agent/src/agent/catalog.rs index 972284abe..6efc0220e 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/catalog.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/catalog.rs @@ -135,9 +135,10 @@ async fn write_model_catalog( } /// The catalog's chat models a task may pick, in catalog order. Aliases are -/// left out: the server may retarget them, and a local agent cannot follow a -/// switch in context window or vision mid-task. A model without -/// `available_in_agent` is offered, since older servers do not send it. +/// left out: the server may retarget one, which would change a task's +/// context window or vision without a switch the task records. A model +/// without `available_in_agent` is offered, since older servers do not +/// send it. pub(super) fn selectable_catalog_models(catalog: &maple_sdk::ModelCatalogResponse) -> Vec { let mut seen = HashSet::new(); catalog @@ -159,6 +160,45 @@ pub(super) fn selectable_catalog_models(catalog: &maple_sdk::ModelCatalogRespons .collect() } +/// The names surfaces show for the catalog's models and aliases, with +/// Research's precedence: the short name, then the full one, then the id. +/// IDs remain the keys used for selection and persisted tasks. +pub fn catalog_model_labels( + catalog: &maple_sdk::ModelCatalogResponse, +) -> std::collections::HashMap { + let label = |short: Option<&str>, full: Option<&str>, id: &str| { + short + .filter(|value| !value.trim().is_empty()) + .or_else(|| full.filter(|value| !value.trim().is_empty())) + .unwrap_or(id) + .to_string() + }; + catalog + .data + .iter() + .map(|model| { + ( + model.id.clone(), + label( + model.short_name.as_deref(), + model.display_name.as_deref(), + &model.id, + ), + ) + }) + .chain(catalog.aliases.iter().map(|alias| { + ( + alias.id.clone(), + label( + alias.short_name.as_deref(), + alias.label.as_deref(), + &alias.id, + ), + ) + })) + .collect() +} + /// The model a new task starts on: the user's last pick, then the server's /// suggestion, then the built-in default, then the catalog's first model; /// each only while the catalog offers it. With no catalog to check against, @@ -391,3 +431,35 @@ impl AgentRuntimeHandle { .to_string()) } } + +#[cfg(test)] +mod tests { + use super::catalog_model_labels; + + #[test] + fn model_labels_follow_research_catalog_precedence() { + let catalog = serde_json::from_value(serde_json::json!({ + "object": "list", + "data": [ + {"id": "model-a", "short_name": "Compact", "display_name": "Full name"}, + {"id": "model-b", "display_name": "Server supplied name"}, + {"id": "model-c", "short_name": " ", "display_name": "Fallback name"}, + {"id": "model-d"} + ], + "aliases": [ + {"id": "alias-a", "short_name": "Alias", "label": "Full alias", "target_model": "model-a"}, + {"id": "alias-b", "label": "Automatic", "target_model": "model-b"}, + {"id": "alias-c", "short_name": "", "label": " ", "target_model": "model-c"} + ] + })) + .unwrap(); + let labels = catalog_model_labels(&catalog); + assert_eq!(labels["model-a"], "Compact"); + assert_eq!(labels["model-b"], "Server supplied name"); + assert_eq!(labels["model-c"], "Fallback name"); + assert_eq!(labels["model-d"], "model-d"); + assert_eq!(labels["alias-a"], "Alias"); + assert_eq!(labels["alias-b"], "Automatic"); + assert_eq!(labels["alias-c"], "alias-c"); + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/agent/runs.rs b/apps/maple-agent/crates/maple-agent/src/agent/runs.rs index fd59b85af..74ae8da21 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/runs.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/runs.rs @@ -355,6 +355,17 @@ impl RunEvents { } } + /// Publish the row of an entry Maple appended itself, linked to the row + /// the store keeps for that entry. + fn publish_entry_row(&self, entry_id: &str, item: AgentTimelineItem) { + { + let mut live = lock(&self.order); + live.drain_changes(&self.session_id); + live.note_entry_row(entry_id, &item.id); + } + self.publish(AgentRunEvent::TimelineItem(item)); + } + fn watch_store(&self, store: &Arc, account: crate::store::AccountId) { let mut live = lock(&self.order); live.changes = Some(store.subscribe_changes()); @@ -1536,15 +1547,8 @@ impl AgentRuntime { .store .get(session_id)? .ok_or_else(|| format!("Failed to find Agent task {session_id}"))?; + // A started task may run on another model: it switches from this run. let model_id = request.model.clone().unwrap_or_else(|| self.model.clone()); - if row.message_count > 0 - && let Some(locked) = row.model.as_deref() - && locked != model_id - { - return Err(format!( - "This task is locked to model {locked}. Start a new task to use {model_id}." - )); - } let model = self .pi_model( &model_id, @@ -1552,7 +1556,14 @@ impl AgentRuntime { Some(request.vision_capable), ) .await; - let session = self.task_session(&row, model).await?; + let (session, switched) = self.switch_task_session(&row, model).await?; + if let Some((entry_id, notice)) = switched { + // Surfaces show the task's new model from the start of the run. + if let Some(row) = self.store.get(session_id)? { + events.publish(AgentRunEvent::SessionUpdated(row.summary())); + } + events.publish_entry_row(&entry_id, notice); + } self.failures.clear(session_id); self.wake(session_id)?; // Named from the first message with text: images alone say nothing, diff --git a/apps/maple-agent/crates/maple-agent/src/agent/runs/live.rs b/apps/maple-agent/crates/maple-agent/src/agent/runs/live.rs index 5c6bddd42..ad5e87c2f 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/runs/live.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/runs/live.rs @@ -77,6 +77,21 @@ impl LiveSnapshot { self.sources.push_back(receipt); } + /// A row of an entry Maple appended itself, such as a notice, as that + /// entry's only row. The live row takes the stored row's identity, so a + /// reload during the run shows it once. + pub(super) fn note_entry_row(&mut self, entry_id: &str, row_id: &str) { + let occurrence = self.occurrence(self.group, row_id); + let origin = Origin { + entry_id: entry_id.into(), + item_index: 0, + }; + self.origins.insert(origin.clone(), occurrence); + if let Some(commit) = self.commits.get(&origin).cloned() { + self.persisted.insert(occurrence, commit); + } + } + pub(super) fn end_message(&mut self, message: &SessionMessage) -> Vec { self.completed_source = None; let signature = message_signature(message); diff --git a/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs b/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs index f0f94000d..cd3c660d6 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs @@ -11,6 +11,7 @@ use std::time::{Duration, Instant}; use pi_agent_core::QueueMode; use pi_ai::Model; use pi_coding_agent::resources::Resources; +use pi_coding_agent::session::EntryKind; use pi_coding_agent::settings::{RetrySettings, Settings}; use pi_coding_agent::{AgentSession, AgentSessionOptions, ModelRegistry}; use tokio_util::sync::CancellationToken; @@ -30,9 +31,11 @@ use super::questions::QuestionBroker; use super::runs::{Failures, Runs}; use super::store::{TaskKind, TaskRow, TaskStore}; use super::surface::{AgentSurfaceAccess, SurfaceTask, with_surface_servers}; +use super::switch_compaction::{OfferedModel, switch_compaction_extension}; +use super::timeline::{MAPLE_NOTICE_ENTRY, notice_entry_data, notice_item}; use super::{ - AgentRuntimeStatus, AgentServiceEvent, MapleAgentHostResources, emit_agent_event, login_path, - tools, + AgentRuntimeStatus, AgentServiceEvent, AgentTimelineItem, MapleAgentHostResources, + emit_agent_event, login_path, tools, }; use crate::maple_api::MapleApiSession; @@ -647,6 +650,16 @@ impl AgentRuntime { row: &TaskRow, model: Model, ) -> Result { + Ok(self.switch_task_session(row, model).await?.0) + } + + /// [`Self::task_session`], with the transcript's notice, and the entry + /// that holds it, when a started task switched to another model. + pub(super) async fn switch_task_session( + &self, + row: &TaskRow, + model: Model, + ) -> Result<(AgentSession, Option<(String, AgentTimelineItem)>), String> { // Read before the sessions are locked: the first read asks a shell, // and the resources come from disk. let search_path = login_path::login_search_path().await; @@ -671,12 +684,13 @@ impl AgentRuntime { return Err("Agent task is being unloaded".into()); } *last_used = Instant::now(); - // A new context window or vision flag for the same model, or - // another model for a task that has not started. - if session.model().as_ref() != Some(&model) { - self.models.register_models([model.clone()]); - session.set_model(model).await; - } + // Another model, or a new context window or vision flag for the + // same one. + let switched = if session.model().as_ref() != Some(&model) { + self.switch_model(session, row, model).await + } else { + None + }; // The web switch and the external agents may have changed since // the last run. tools::sync_web_tools(session, row.web_enabled); @@ -684,17 +698,28 @@ impl AgentRuntime { if let Some(resources) = resources { session.set_resources(resources); } - return Ok(session.clone()); + return Ok((session.clone(), switched)); } - // Registered first, so a resumed session finds the model it names. - self.models.register_models([model.clone()]); let store = Arc::clone(&self.store); let (id, cwd) = (row.id.clone(), row.project_root.clone()); let manager = tokio::task::spawn_blocking(move || store.open_session(&id, &cwd)) .await .map_err(|error| format!("Failed to read the Agent task: {error}"))??; + // A started task opens on the model it last used and then switches, + // as a loaded one does, so the switch is recorded and announced. + let stored = manager + .projection() + .model + .map(|(_, id)| id) + .filter(|id| row.message_count > 0 && *id != model.id); + let initial = match &stored { + Some(id) => self.pi_model(id, None, None).await, + None => model.clone(), + }; + // Registered first, so a resumed session finds the model it names. + self.models.register_models([initial.clone()]); // A model that cannot see images gets them described. - let describe = !model.supports_images(); + let describe = !initial.supports_images(); let attachments = Arc::new(super::attachments::AgentAttachmentStore::new( Arc::clone(self.store.store()), self.store.account(), @@ -702,7 +727,7 @@ impl AgentRuntime { let task_models = self.attachment_models(&row.id); let mut options = AgentSessionOptions::new(&row.project_root, APP_NAME, manager, task_models); - options.model = Some(model); + options.model = Some(initial); options.settings = self.settings(); options.resources = resources.unwrap_or_default(); @@ -747,9 +772,17 @@ impl AgentRuntime { options.extensions.push(mcp.extension()); let task_cua = TaskCua::new(&self.account_scope, &row.id, self.models.clone()); options.extensions.push(task_cua.extension()); + // After a switch to a smaller window, the previous model summarizes. + options + .extensions + .push(switch_compaction_extension(self.offered_model())); let session = AgentSession::new(options) .await .map_err(|error| format!("Failed to start the Agent task: {error}"))?; + let switched = match stored { + Some(_) => self.switch_model(&session, row, model).await, + None => None, + }; sync_external_agent_tools(&session, &agents, providers); sessions.insert( row.id.clone(), @@ -764,7 +797,105 @@ impl AgentRuntime { ); drop(sessions); self.evict_idle_tasks(&row.id).await; - Ok(session) + Ok((session, switched)) + } + + /// Move a session to `model`, as Pi's `setModel` does. A started task + /// that switched to another model records a notice in its transcript, + /// which is returned with its entry to show live. + /// + /// Pi's switch takes the thinking default of the session's settings, + /// read when the session was built. A switch to another model takes the + /// account's default as it is now, which the composer previews. + async fn switch_model( + &self, + session: &AgentSession, + row: &TaskRow, + model: Model, + ) -> Option<(String, AgentTimelineItem)> { + let changed = session + .model() + .is_none_or(|previous| previous.id != model.id); + let model_id = model.id.clone(); + self.models.register_models([model.clone()]); + session.set_model(model).await; + if !changed { + return None; + } + if let Some(level) = self.thinking_default() { + session.set_thinking_level(level).await; + } + if row.message_count == 0 { + return None; + } + let text = format!("Switched to {}", self.model_label(&model_id).await); + let id = format!("model-switch-{}", pi_ai::now_ms()); + session + .extension_context() + .append_entry(MAPLE_NOTICE_ENTRY, Some(notice_entry_data(&id, &text))); + let entry_id = session.with_session(|manager| { + manager + .branch() + .iter() + .rev() + .find_map(|entry| match &entry.kind { + EntryKind::Custom { + custom_type, + data: Some(data), + } if custom_type == MAPLE_NOTICE_ENTRY + && data.get("id").and_then(serde_json::Value::as_str) + == Some(id.as_str()) => + { + Some(entry.id.clone()) + } + _ => None, + }) + })?; + Some(( + entry_id, + notice_item(id, "Agent notice", &text, pi_ai::now_ms()), + )) + } + + /// A model as the account's catalog offers it to tasks: `None` when the + /// catalog does not offer it, or cannot be read. + fn offered_model(&self) -> OfferedModel { + let api = Arc::clone(&self.api); + Arc::new(move |model_id: String| { + let api = Arc::clone(&api); + Box::pin(async move { + let catalog = api.model_catalog().await.ok()?; + if !super::catalog::selectable_catalog_models(&catalog).contains(&model_id) { + return None; + } + let entry = catalog_entry(&catalog, &model_id); + Some(maple_model( + &model_id, + entry.as_ref(), + context_limit_override(), + )) + }) + }) + } + + /// The name the catalog shows for a model, or its id. + async fn model_label(&self, model_id: &str) -> String { + self.api + .model_catalog() + .await + .ok() + .and_then(|catalog| super::catalog_model_labels(&catalog).remove(model_id)) + .unwrap_or_else(|| model_id.to_string()) + } + + /// The account's saved thinking default, if it has one. + fn thinking_default(&self) -> Option { + super::config::load_agent_config_inner(&self.host.paths, &self.user_id) + .map(|config| config.default_thinking_level) + .unwrap_or_else(|error| { + log::warn!("Failed to read the thinking default: {error}"); + None + }) } /// The skills, prompt templates and instruction files of a task in @@ -805,15 +936,7 @@ impl AgentRuntime { http_idle_timeout_ms: Some(0), steering_mode: QueueMode::All, follow_up_mode: QueueMode::All, - default_thinking_level: super::config::load_agent_config_inner( - &self.host.paths, - &self.user_id, - ) - .map(|config| config.default_thinking_level) - .unwrap_or_else(|error| { - log::warn!("Failed to read the thinking default: {error}"); - None - }), + default_thinking_level: self.thinking_default(), append_system_prompt: Some(appended).filter(|text| !text.is_empty()), ..Settings::default() } diff --git a/apps/maple-agent/crates/maple-agent/src/agent/store/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/store/tests.rs index 402b8e9a9..27e0c51ae 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/store/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/store/tests.rs @@ -222,23 +222,31 @@ fn metadata_compare_and_swap_preserves_a_concurrent_user_rename() { assert!(saved.title_user_set); } +/// A started task takes a selection for its next run; the session's next +/// append then records the model the session has. #[test] -fn model_selection_is_locked_by_the_same_transaction_that_observes_first_message() { +fn a_started_task_takes_a_model_selection_until_its_session_records_one() { let (_directory, store) = store(); store.insert(&row(1, "one", 1)).unwrap(); store .update(&id(1), |row| row.model = Some("initial-selection".into())) .unwrap(); - let result = store.update(&id(1), |row| { - let mut session = store.open_session(&id(1), &root("one")).unwrap(); - session.append_model_change("maple", "initial-selection"); - session.append_message(user("first message")); - row.model = Some("racing-selection".into()); - }); - assert!(result.is_err()); + let mut session = store.open_session(&id(1), &root("one")).unwrap(); + session.append_model_change("maple", "initial-selection"); + session.append_message(user("first message")); + + let selected = store + .update(&id(1), |row| row.model = Some("next-selection".into())) + .unwrap() + .unwrap(); + assert_eq!(selected.model.as_deref(), Some("next-selection")); + assert_eq!(selected.message_count, 1); + + session.append_model_change("maple", "next-selection"); + session.append_message(user("second message")); let saved = store.get(&id(1)).unwrap().unwrap(); - assert_eq!(saved.model.as_deref(), Some("initial-selection")); - assert_eq!(saved.message_count, 1); + assert_eq!(saved.model.as_deref(), Some("next-selection")); + assert_eq!(saved.message_count, 2); } #[test] diff --git a/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction.rs b/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction.rs new file mode 100644 index 000000000..ab325adc1 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction.rs @@ -0,0 +1,254 @@ +//! Compaction after a task switched to a model with a smaller context +//! window, as a Pi extension. +//! +//! Pi compacts with the session's model, in one request. After a switch to +//! a smaller window, the history to summarize was written for the previous +//! model and may not fit the new one. The model that wrote the last reply +//! held all of it, so it writes the summary instead. It is read from the +//! stored reply, so this holds across restarts. When the catalog no longer +//! offers that model, or its summary fails, the new model summarizes the +//! history trimmed to its own window, oldest messages first. +//! +//! Every other compaction is Pi's own. So is this one when the new model +//! can take the whole history, or when both attempts fail: the extension +//! then returns no summary. + +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use pi_ai::{Message, Model, StreamOptions}; +use pi_coding_agent::compaction::{ + CompactionPreparation, CompactionResult, Summarizer, compact, serialize_conversation, +}; +use pi_coding_agent::extensions::{ + BeforeCompactResult, Extension, ExtensionContext, SessionBeforeCompact, extension, +}; +use pi_coding_agent::messages::convert_to_llm; +use pi_coding_agent::session::{EntryKind, SessionEntry}; +use pi_coding_agent::{ModelRegistry, SessionMessage}; + +/// A model as the account's catalog offers it to tasks, by id, or `None` +/// when the catalog does not offer it. +pub(crate) type OfferedModel = + Arc BoxFuture<'static, Option> + Send + Sync>; + +/// Room for a summary request's system prompt, instructions and framing. +const PROMPT_TOKENS: u64 = 2_000; + +pub(crate) fn switch_compaction_extension(offered: OfferedModel) -> Arc { + extension("maple-switch-compaction", move |api| { + let offered = Arc::clone(&offered); + api.on( + move |event: SessionBeforeCompact, context: ExtensionContext| { + let offered = Arc::clone(&offered); + async move { Ok(compact_after_switch(event, &context, &offered).await) } + }, + ); + }) +} + +async fn compact_after_switch( + event: SessionBeforeCompact, + context: &ExtensionContext, + offered: &OfferedModel, +) -> BeforeCompactResult { + let (Some(current), Some(models)) = (context.model(), context.models()) else { + return BeforeCompactResult::default(); + }; + let Some(previous_id) = switched_from(&event.branch_entries, ¤t) else { + return BeforeCompactResult::default(); + }; + let previous = offered(previous_id.clone()).await; + // A switch to a window at least as large leaves nothing to fit. + if previous + .as_ref() + .is_some_and(|previous| previous.context_window <= current.context_window) + { + return BeforeCompactResult::default(); + } + let summaries = Summaries { + models, + context, + event: &event, + }; + match &previous { + Some(previous) => match summaries.write(previous, &event.preparation).await { + Ok(result) => { + log::info!( + "Compacted after a switch from {} to {}: summarized by {}", + previous.id, + current.id, + previous.id + ); + return summarized(result); + } + Err(error) => log::warn!( + "Compaction after a switch to {}: {} could not summarize: {error}", + current.id, + previous.id + ), + }, + None => log::info!( + "Compaction after a switch to {}: {previous_id} is no longer offered", + current.id + ), + } + if event.signal.is_cancelled() { + return cancelled(); + } + let Some(trimmed) = fit_preparation( + &event.preparation, + ¤t, + event.custom_instructions.as_deref(), + ) else { + // The new model takes all of it, as Pi's own summary does. + return BeforeCompactResult::default(); + }; + match summaries.write(¤t, &trimmed).await { + Ok(result) => { + log::info!( + "Compacted after a switch to {}: summarized by {} from {} of {} messages", + current.id, + current.id, + trimmed.messages_to_summarize.len() + trimmed.turn_prefix_messages.len(), + event.preparation.messages_to_summarize.len() + + event.preparation.turn_prefix_messages.len() + ); + summarized(result) + } + Err(_) if event.signal.is_cancelled() => cancelled(), + Err(error) => { + log::warn!( + "Compaction after a switch to {}: the trimmed summary failed: {error}", + current.id + ); + BeforeCompactResult::default() + } + } +} + +fn summarized(result: CompactionResult) -> BeforeCompactResult { + BeforeCompactResult { + compaction: Some(result), + ..BeforeCompactResult::default() + } +} + +fn cancelled() -> BeforeCompactResult { + BeforeCompactResult { + cancel: true, + ..BeforeCompactResult::default() + } +} + +/// The model of the branch's last complete reply, when it is not `current`. +fn switched_from(branch: &[SessionEntry], current: &Model) -> Option { + let reply = branch.iter().rev().find_map(|entry| match &entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::Assistant(reply)), + } if !reply.is_failure() => Some(reply), + _ => None, + })?; + (reply.provider != current.provider || reply.model != current.id).then(|| reply.model.clone()) +} + +/// Summary requests made as Pi makes them, on a model of choice. +struct Summaries<'a> { + models: ModelRegistry, + context: &'a ExtensionContext, + event: &'a SessionBeforeCompact, +} + +impl Summaries<'_> { + async fn write( + &self, + model: &Model, + preparation: &CompactionPreparation, + ) -> Result { + let settings = self.context.settings(); + let stream_fn = self.models.stream_fn(); + let summarizer = Summarizer { + model, + stream_fn: &*stream_fn, + options: StreamOptions { + api_key: self.models.api_key(&model.provider).await, + reasoning: Some(self.context.thinking_level()), + session_id: self.context.session_id(), + cancel: self.event.signal.clone(), + ..StreamOptions::default() + }, + max_retries: if settings.retry.enabled { + settings.retry.max_retries + } else { + 0 + }, + retry_base_ms: settings.retry.base_delay_ms, + retry_max_ms: settings.retry.max_agent_delay_ms, + on_retry: None, + }; + compact( + preparation, + &summarizer, + self.event.custom_instructions.as_deref(), + ) + .await + } +} + +/// `preparation` with the oldest messages of each summary request left out +/// until the request fits `model`'s window with room for the summary, or +/// `None` when it fits whole. +fn fit_preparation( + preparation: &CompactionPreparation, + model: &Model, + instructions: Option<&str>, +) -> Option { + let budget = model + .context_window + .saturating_sub(preparation.settings.reserve_tokens) + .saturating_sub(PROMPT_TOKENS) + .saturating_sub(text_tokens(instructions.unwrap_or_default())); + let history_budget = budget.saturating_sub(text_tokens( + preparation.previous_summary.as_deref().unwrap_or_default(), + )); + let messages_to_summarize = newest_within(&preparation.messages_to_summarize, history_budget); + let turn_prefix_messages = newest_within(&preparation.turn_prefix_messages, budget); + let whole = messages_to_summarize.len() == preparation.messages_to_summarize.len() + && turn_prefix_messages.len() == preparation.turn_prefix_messages.len(); + (!whole).then(|| CompactionPreparation { + messages_to_summarize, + turn_prefix_messages, + ..preparation.clone() + }) +} + +/// The newest of `messages` whose text, as a summary request shows it, fits +/// `budget` tokens. +fn newest_within(messages: &[SessionMessage], budget: u64) -> Vec { + let mut total = 0u64; + let start = messages + .iter() + .enumerate() + .rev() + .find(|(_, message)| { + total = total.saturating_add(request_tokens(message)); + total > budget + }) + .map_or(0, |(index, _)| index + 1); + messages[start..].to_vec() +} + +/// A message's share of a summary request: its rendered text and the +/// separator before it, estimated as Pi's compaction estimates. +fn request_tokens(message: &SessionMessage) -> u64 { + text_tokens(&serialize_conversation(&convert_to_llm( + std::slice::from_ref(message), + ))) + 1 +} + +fn text_tokens(text: &str) -> u64 { + text.chars().count().div_ceil(4) as u64 +} + +#[cfg(test)] +mod tests; diff --git a/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction/tests.rs new file mode 100644 index 000000000..3985091aa --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/agent/switch_compaction/tests.rs @@ -0,0 +1,302 @@ +//! The extension's decisions, on a Pi session of its own with a scripted +//! provider: which compactions it leaves to Pi, which model summarizes the +//! others, and what the trimmed history keeps. + +use std::sync::{Arc, Mutex}; + +use pi_ai::faux::{FauxProvider, FauxRequest, faux_message}; +use pi_ai::{AssistantContent, Message, Model, UserMessage}; +use pi_coding_agent::compaction::CompactionSettings; +use pi_coding_agent::session::SessionManager; +use pi_coding_agent::{AgentSession, AgentSessionOptions, ModelRegistry, StaticKeys}; + +use super::*; + +const OLDEST: &str = "oldest-marker"; +const KEPT: &str = "kept-marker"; + +fn model(id: &str, context_window: u64) -> Model { + Model { + id: id.to_string(), + context_window, + ..FauxProvider::default_model() + } +} + +fn big() -> Model { + model("big", 200_000) +} + +fn small() -> Model { + model("small", 40_000) +} + +/// A user message of about `tokens` tokens, as compaction estimates them. +fn user(marker: &str, tokens: usize) -> SessionMessage { + SessionMessage::Llm(Message::User(UserMessage::text(format!( + "{marker} {}", + "x".repeat(tokens * 4) + )))) +} + +fn reply(model: &Model, text: &str) -> SessionMessage { + let mut reply = faux_message(vec![AssistantContent::text(text)]); + reply.provider = model.provider.clone(); + reply.model = model.id.clone(); + SessionMessage::Llm(Message::Assistant(reply)) +} + +/// A history the previous model wrote: its oldest message is about 30k +/// tokens, which the small window cannot summarize, and the newest 25k +/// stay after the cut. +fn history(writer: &Model) -> Vec { + vec![ + user(OLDEST, 30_000), + reply(writer, KEPT), + user("recent", 25_000), + reply(writer, "Recent answer."), + ] +} + +/// What the catalog offers: `offered`, whatever the id asked for. +fn offering(offered: Option) -> OfferedModel { + Arc::new(move |_id| { + let offered = offered.clone(); + Box::pin(async move { offered }) + }) +} + +async fn session_on( + current: Model, + history: Vec, + offered: OfferedModel, + faux: &FauxProvider, +) -> AgentSession { + let registry = ModelRegistry::new(Arc::new(StaticKeys::default())); + registry.register_api("faux", Arc::new(faux.clone())); + let mut manager = SessionManager::in_memory("/project"); + for message in history { + manager.append_message(message); + } + let mut options = AgentSessionOptions::new("/project", "Maple", manager, registry); + options.model = Some(current); + options.settings.retry.base_delay_ms = 1; + options + .extensions + .push(switch_compaction_extension(offered)); + AgentSession::new(options).await.unwrap() +} + +/// Compact the session by hand, as `/compact` does. Returns whether the +/// extension wrote the summary, and the summary requests in order. +async fn compact_now(session: &AgentSession, faux: &FauxProvider) -> (bool, Vec) { + session.compact(None).await.unwrap(); + let from_extension = session.with_session(|manager| { + manager + .branch() + .iter() + .rev() + .find_map(|entry| match &entry.kind { + EntryKind::Compaction { from_extension, .. } => Some(*from_extension), + _ => None, + }) + }); + (from_extension.expect("compacted"), faux.requests()) +} + +fn request_text(request: &FauxRequest) -> String { + request + .context + .messages + .iter() + .filter_map(|message| match message { + Message::User(user) => Some(pi_ai::content_text(&user.content)), + _ => None, + }) + .collect() +} + +fn models(requests: &[FauxRequest]) -> Vec<&str> { + requests + .iter() + .map(|request| request.model.id.as_str()) + .collect() +} + +#[tokio::test] +async fn a_compaction_on_the_same_model_is_pis_own() { + let faux = FauxProvider::new(); + faux.push_text("Summary."); + let session = session_on(small(), history(&small()), offering(Some(big())), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(!from_extension); + assert_eq!(models(&requests), ["small"]); +} + +#[tokio::test] +async fn a_switch_to_a_larger_window_is_pis_own() { + let faux = FauxProvider::new(); + faux.push_text("Summary."); + let session = session_on(big(), history(&small()), offering(Some(small())), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(!from_extension); + assert_eq!(models(&requests), ["big"]); +} + +#[tokio::test] +async fn a_switch_to_a_smaller_window_is_summarized_by_the_previous_model() { + let faux = FauxProvider::new(); + faux.push_text("Summary by the previous model."); + let asked = Arc::new(Mutex::new(Vec::new())); + let offered: OfferedModel = { + let asked = Arc::clone(&asked); + Arc::new(move |id| { + asked.lock().unwrap().push(id); + Box::pin(async { Some(big()) }) + }) + }; + let session = session_on(small(), history(&big()), offered, &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(from_extension); + assert_eq!(*asked.lock().unwrap(), ["big"]); + assert_eq!(models(&requests), ["big"]); + assert!(request_text(&requests[0]).contains(OLDEST)); + let summary = session.with_session(|manager| { + manager.branch().iter().find_map(|entry| match &entry.kind { + EntryKind::Compaction { summary, .. } => Some(summary.clone()), + _ => None, + }) + }); + assert!( + summary + .unwrap() + .starts_with("Summary by the previous model."), + "the previous model's summary is recorded" + ); +} + +#[tokio::test] +async fn the_new_model_summarizes_trimmed_history_when_the_previous_one_fails() { + let faux = FauxProvider::new(); + faux.push_error("The Maple inference endpoint was not found"); + faux.push_text("Summary of what fits."); + let session = session_on(small(), history(&big()), offering(Some(big())), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(from_extension); + assert_eq!(models(&requests), ["big", "small"]); + let trimmed = request_text(&requests[1]); + assert!(!trimmed.contains(OLDEST), "the oldest message is left out"); + assert!(trimmed.contains(KEPT), "newer messages stay"); +} + +#[tokio::test] +async fn the_new_model_summarizes_trimmed_history_when_the_previous_one_is_gone() { + let faux = FauxProvider::new(); + faux.push_text("Summary of what fits."); + let session = session_on(small(), history(&big()), offering(None), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(from_extension); + assert_eq!(models(&requests), ["small"]); + assert!(!request_text(&requests[0]).contains(OLDEST)); +} + +/// A history the new model takes whole needs nothing from the extension. +#[tokio::test] +async fn pi_summarizes_when_the_new_model_takes_the_whole_history() { + let faux = FauxProvider::new(); + faux.push_text("Pi's own summary."); + let history = vec![ + user("short", 1_000), + reply(&big(), KEPT), + user("recent", 25_000), + reply(&big(), "Recent answer."), + ]; + let session = session_on(small(), history, offering(None), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(!from_extension); + assert_eq!(models(&requests), ["small"]); +} + +#[tokio::test] +async fn pi_summarizes_when_both_attempts_fail() { + let faux = FauxProvider::new(); + faux.push_error("The Maple inference endpoint was not found"); + faux.push_error("The Maple inference endpoint was not found"); + faux.push_text("Pi's own summary."); + let session = session_on(small(), history(&big()), offering(Some(big())), &faux).await; + let (from_extension, requests) = compact_now(&session, &faux).await; + assert!(!from_extension); + assert_eq!(models(&requests), ["big", "small", "small"]); + assert!( + request_text(&requests[2]).contains(OLDEST), + "Pi sends it all" + ); +} + +#[test] +fn the_switch_is_read_from_the_last_complete_reply() { + let mut manager = SessionManager::in_memory("/project"); + let current = small(); + manager.append_message(user("first", 1)); + assert_eq!(switched_from(&branch(&manager), ¤t), None); + manager.append_message(reply(&big(), "From the previous model.")); + assert_eq!( + switched_from(&branch(&manager), ¤t).as_deref(), + Some("big") + ); + // A failed reply on the new model wrote nothing. + let mut failed = faux_message(Vec::new()); + failed.provider = current.provider.clone(); + failed.model = current.id.clone(); + failed.stop_reason = pi_ai::StopReason::Error; + manager.append_message(SessionMessage::Llm(Message::Assistant(failed))); + assert_eq!( + switched_from(&branch(&manager), ¤t).as_deref(), + Some("big") + ); + manager.append_message(reply(¤t, "From the new model.")); + assert_eq!(switched_from(&branch(&manager), ¤t), None); +} + +fn branch(manager: &SessionManager) -> Vec { + manager.branch().into_iter().cloned().collect() +} + +#[test] +fn trimmed_history_keeps_the_newest_messages_that_fit() { + let preparation = CompactionPreparation { + first_kept_entry_id: "kept".into(), + messages_to_summarize: vec![ + user(OLDEST, 30_000), + reply(&big(), KEPT), + user("newer", 5_000), + ], + turn_prefix_messages: vec![user("prefix", 1_000)], + tokens_before: 60_000, + previous_summary: Some("y".repeat(4 * 1_000)), + file_operations: Default::default(), + settings: CompactionSettings::default(), + }; + let window = small().context_window; + let fitted = fit_preparation(&preparation, &small(), None).unwrap(); + assert_eq!(fitted.messages_to_summarize.len(), 2); + assert_eq!( + fitted.turn_prefix_messages, + preparation.turn_prefix_messages + ); + assert_eq!(fitted.first_kept_entry_id, "kept"); + assert_eq!(fitted.previous_summary, preparation.previous_summary); + let request: u64 = fitted + .messages_to_summarize + .iter() + .map(request_tokens) + .sum::() + + text_tokens(fitted.previous_summary.as_deref().unwrap()) + + PROMPT_TOKENS; + assert!( + request <= window - CompactionSettings::default().reserve_tokens, + "{request}" + ); + // Everything fits a window large enough. + assert_eq!(fit_preparation(&preparation, &big(), None), None); +} diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs b/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs index 9dd819a99..34b2d2e1a 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs @@ -473,8 +473,8 @@ impl AgentRuntimeHandle { .collect()) } - /// A task's summary, transcript and queue. - /// Change a draft task's model at the same boundary that admits its first run. + /// Select the model of a task's next run, at the same boundary that + /// admits runs. A started task switches when that run starts. pub async fn set_task_model( &self, session_id: &str, @@ -483,10 +483,6 @@ impl AgentRuntimeHandle { self.verify_generation().await?; let runtime = self.runtime().await?; let summary = runtime.runs.while_idle(session_id, || { - let row = runtime.store.get(session_id)?.ok_or("Task was not found")?; - if row.message_count > 0 && row.model.as_deref() != Some(model) { - return Err("Maple tasks are model-locked after their first message".into()); - } runtime .store .update(session_id, |row| row.model = Some(model.into()))? @@ -504,6 +500,7 @@ impl AgentRuntimeHandle { Ok(summary) } + /// A task's summary, transcript and queue. pub async fn load_session(&self, session_id: String) -> Result { self.ensure_task_history_ready(&session_id).await?; let runtime = self.current_runtime().await?; diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/tests.rs index ed7e79ecf..7b55629e6 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tests.rs @@ -1994,9 +1994,64 @@ async fn a_failed_reply_fails_the_run_and_shows_its_error() { assert_eq!(error.title.as_deref(), Some("Credits exhausted")); } +/// The models a task's session recorded switching to, in order. +fn model_changes(session: &pi_coding_agent::AgentSession) -> Vec { + session.with_session(|manager| { + manager + .branch() + .iter() + .filter_map(|entry| match &entry.kind { + pi_coding_agent::session::EntryKind::ModelChange { model_id, .. } => { + Some(model_id.clone()) + } + _ => None, + }) + .collect() + }) +} + +/// The switch notice a run showed live. It carries the identity of its +/// stored row, so a reload during the run (after a compaction, for one) +/// shows it once. +fn live_switch_notice(harness: &Harness, run_id: &str) -> AgentTimelineItem { + let notice = harness + .recorder + .run_events(run_id) + .into_iter() + .find_map(|event| match event { + AgentRunEvent::TimelineItem(item) if item.id.starts_with("model-switch-") => Some(item), + _ => None, + }) + .expect("the run shows the switch"); + assert!(notice.persisted.is_some(), "{notice:?}"); + notice +} + +/// The texts of a timeline's notices and user messages, in order. +fn notices_and_prompts(timeline: &[AgentTimelineItem]) -> Vec { + timeline + .iter() + .filter(|row| row.item_type == "system" || row.role.as_deref() == Some("user")) + .filter_map(|row| row.text.clone()) + .collect() +} + +/// A started task switches models with its next run, loaded or not: the +/// run goes to the new model with the whole history, the task's model +/// follows from the start of the run, and the transcript says where it +/// switched, live and once stored. #[tokio::test] -async fn a_run_wakes_a_settled_task_and_a_started_task_keeps_its_model() { - let harness = Harness::new().await; +async fn a_run_wakes_a_settled_task_and_a_started_task_switches_models() { + let mut harness = Harness::new().await; + let catalog = json!({ + "object": "list", + "data": [ + {"id": "glm-5-3", "short_name": "GLM 5.3"}, + {"id": "kimi-k3", "short_name": "Kimi K3"}, + ], + "aliases": [], + }); + harness.serve_catalog(catalog.clone()).await; harness.faux.push_text("Hello."); let task = harness.create_task().await; harness @@ -2009,15 +2064,57 @@ async fn a_run_wakes_a_settled_task_and_a_started_task_keeps_its_model() { let listed = harness.handle.list_sessions(None).await.unwrap(); assert_eq!(listed[0].state, AgentTaskState::Active); + harness.faux.push_text("Hello from Kimi."); let mut other_model = harness.request(&task, "Again"); other_model.model = Some("kimi-k3".to_string()); - let error = harness - .handle - .send_message(other_model) - .await - .err() - .unwrap(); - assert!(error.contains("locked to model glm-5-3"), "{error}"); + let mut run = harness.handle.send_message(other_model).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let request = harness.faux.requests().pop().unwrap(); + assert_eq!(request.model.id, "kimi-k3"); + assert_eq!(user_texts(&request), ["Hi", "Again"]); + let switched = harness + .recorder + .run_events(&run.run_id) + .into_iter() + .find_map(|event| match event { + AgentRunEvent::SessionUpdated(summary) => summary.model, + _ => None, + }); + assert_eq!(switched.as_deref(), Some("kimi-k3")); + let live = notices_and_prompts(&harness.recorder.live_rows(&run.run_id)); + assert_eq!(live, ["Switched to Kimi K3", "Again"]); + live_switch_notice(&harness, &run.run_id); + let stored = harness.handle.load_session(task.clone()).await.unwrap(); + assert_eq!( + notices_and_prompts(&stored.timeline), + ["Hi", "Switched to Kimi K3", "Again"] + ); + + // After a restart the task opens on Kimi and switches back the same way. + harness.restart().await; + harness.start().await; + harness.serve_catalog(catalog).await; + harness.faux.push_text("Hello from GLM."); + let mut run = harness.send(&task, "Once more").await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let request = harness.faux.requests().pop().unwrap(); + assert_eq!(request.model.id, "glm-5-3"); + assert_eq!(user_texts(&request), ["Hi", "Again", "Once more"]); + let session = harness.runtime().await.loaded_session(&task).await.unwrap(); + assert_eq!(model_changes(&session), ["glm-5-3", "kimi-k3", "glm-5-3"]); + let listed = harness.handle.list_sessions(None).await.unwrap(); + assert_eq!(listed[0].model.as_deref(), Some("glm-5-3")); + let stored = harness.handle.load_session(task).await.unwrap(); + assert_eq!( + notices_and_prompts(&stored.timeline), + [ + "Hi", + "Switched to Kimi K3", + "Again", + "Switched to GLM 5.3", + "Once more" + ] + ); } #[tokio::test] @@ -3288,6 +3385,230 @@ async fn thinking_picks_persist_in_pi_and_default_new_tasks_after_restart() { ); } +/// The start of the oldest message of [`long_task`], which a small window +/// cannot summarize. +const OLDEST_MESSAGE: &str = "Oldest request"; + +/// A runtime whose catalog offers GLM 5.3 with a large window and Kimi K3 +/// with a small one, and a task on GLM long enough to compact on Kimi: +/// its oldest message is about 30k tokens and its newest 25k stay. +async fn long_task() -> (Harness, String) { + let harness = Harness::new().await; + harness + .serve_catalog(json!({ + "object": "list", + "data": [ + {"id": "glm-5-3", "short_name": "GLM 5.3", "context_window": 200_000}, + {"id": "kimi-k3", "short_name": "Kimi K3", "context_window": 40_000}, + ], + "aliases": [], + })) + .await; + let task = harness.create_task().await; + for (text, reply) in [ + ( + format!("{OLDEST_MESSAGE} {}", "x".repeat(120_000)), + "Noted.", + ), + ( + format!("Recent request {}", "y".repeat(100_000)), + "Noted again.", + ), + ] { + harness.faux.push_text(reply); + let mut run = harness.send(&task, &text).await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + } + (harness, task) +} + +fn on_kimi(harness: &Harness, task: &str, text: &str) -> AgentSendMessageRequest { + AgentSendMessageRequest { + model: Some("kimi-k3".to_string()), + context_limit: Some(40_000), + ..harness.request(task, text) + } +} + +/// What a request asked to summarize, or `None` for a task's turn. +fn summarized(request: &pi_ai::faux::FauxRequest) -> Option { + let system = request + .context + .messages + .iter() + .find_map(|message| match message { + pi_ai::Message::System(system) => Some(system.content.clone()), + _ => None, + })?; + (system == pi_coding_agent::compaction::SUMMARIZATION_SYSTEM_PROMPT) + .then(|| last_user_text(request)) +} + +/// A switch to a smaller window compacts before the next prompt with the +/// model that wrote the history, and the run continues on the new one. +#[tokio::test] +async fn a_switch_to_a_smaller_window_compacts_with_the_previous_model() { + let (harness, task) = long_task().await; + harness.faux.push_text("Summary by GLM."); + harness.faux.push_text("Answer from Kimi."); + let mut run = harness + .handle + .send_message(on_kimi(&harness, &task, "Go on")) + .await + .unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let notice = live_switch_notice(&harness, &run.run_id); + assert_eq!(notice.text.as_deref(), Some("Switched to Kimi K3")); + + let requests = harness.faux.requests(); + let [summary, answer] = &requests[2..] else { + panic!("a summary and a turn: {}", requests.len()) + }; + assert_eq!(summary.model.id, "glm-5-3"); + assert!(summarized(summary).unwrap().contains(OLDEST_MESSAGE)); + assert_eq!(answer.model.id, "kimi-k3"); + assert_eq!(summarized(answer), None); + let context = user_texts(answer).join("\n"); + assert!(context.contains("Summary by GLM."), "{context:.200}"); + assert!(!context.contains(OLDEST_MESSAGE)); + + let detail = harness.handle.load_session(task).await.unwrap(); + let shown = notices_and_prompts(&detail.timeline); + assert_eq!( + shown[shown.len() - 3..], + [ + "Switched to Kimi K3", + "Earlier messages were summarized to make room in the context.", + "Go on" + ] + ); + assert_eq!( + detail.timeline.last().unwrap().text.as_deref(), + Some("Answer from Kimi.") + ); +} + +/// When the previous model cannot summarize, the new one summarizes the +/// history trimmed to its window, and the run continues. +#[tokio::test] +async fn a_failed_previous_summary_falls_back_to_trimmed_history_on_the_new_model() { + let (harness, task) = long_task().await; + harness + .faux + .push_error("The Maple inference endpoint was not found"); + harness.faux.push_text("Summary by Kimi."); + harness.faux.push_text("Answer from Kimi."); + let mut run = harness + .handle + .send_message(on_kimi(&harness, &task, "Go on")) + .await + .unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + + let requests = harness.faux.requests(); + let [failed, trimmed, answer] = &requests[2..] else { + panic!("two summaries and a turn: {}", requests.len()) + }; + assert_eq!(failed.model.id, "glm-5-3"); + assert!(summarized(failed).is_some()); + assert_eq!(trimmed.model.id, "kimi-k3"); + let input = summarized(trimmed).unwrap(); + assert!( + !input.contains(OLDEST_MESSAGE), + "the oldest message is left out" + ); + assert!(input.contains("Noted."), "newer messages stay"); + assert_eq!(answer.model.id, "kimi-k3"); + assert!(user_texts(answer).join("\n").contains("Summary by Kimi.")); +} + +/// A started task previews the model it will switch to with the level the +/// switch gives it, keeps a pick for that model off its current one, and +/// takes the pick when its next run switches. +#[tokio::test] +async fn a_started_task_previews_and_takes_thinking_for_its_next_model() { + use ThinkingLevel::{High, Low, Max}; + let harness = Harness::new().await; + harness + .serve_catalog(json!({"object":"list","data":[ + {"id":"glm-5-3", "reasoning":{"mandatory":true,"supported_efforts":["max","high","low"]}}, + {"id":"kimi-k3", "reasoning":{"supported_efforts":["high","low"]}} + ],"aliases":[]})) + .await; + let task = harness.create_task().await; + harness + .handle + .set_thinking_level(Some(&task), "glm-5-3", Max) + .await + .unwrap(); + harness.faux.push_text("Done."); + let mut run = harness.send(&task, "Hello").await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + + let state = async |model: &str| { + harness + .handle + .thinking_state(Some(&task), model) + .await + .unwrap() + }; + assert_eq!( + state("kimi-k3").await, + AgentThinkingState { + levels: vec![Low, High], + current: High, + } + ); + assert_eq!(state("glm-5-3").await.current, Max); + + let picked = harness + .handle + .set_thinking_level(Some(&task), "kimi-k3", Low) + .await + .unwrap(); + assert_eq!(picked.current, Low); + let session = harness.runtime().await.loaded_session(&task).await.unwrap(); + assert_eq!(session.model().unwrap().id, "glm-5-3"); + assert_eq!(session.thinking_level(), Max); + + harness.faux.push_text("Switched."); + let mut kimi = harness.request(&task, "Go on"); + kimi.model = Some("kimi-k3".to_string()); + let mut run = harness.handle.send_message(kimi).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + assert_eq!(session.model().unwrap().id, "kimi-k3"); + assert_eq!(session.thinking_level(), Low); + assert_eq!(state("kimi-k3").await.current, Low); +} + +/// An ACP or other surface selection on a started task is stored for its +/// next run, and the run switches to it. +#[tokio::test] +async fn a_started_task_takes_a_model_selection_for_its_next_run() { + let harness = Harness::new().await; + harness.faux.push_text("Hello."); + let task = harness.create_task().await; + let mut run = harness.send(&task, "Hi").await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + + let selected = harness + .handle + .set_task_model(&task, "kimi-k3") + .await + .unwrap(); + assert_eq!(selected.model.as_deref(), Some("kimi-k3")); + assert_eq!(selected.message_count, 2); + + harness.faux.push_text("Hello from Kimi."); + let mut kimi = harness.request(&task, "Again"); + kimi.model = selected.model; + let mut run = harness.handle.send_message(kimi).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + assert_eq!(harness.faux.requests().pop().unwrap().model.id, "kimi-k3"); + let listed = harness.handle.list_sessions(None).await.unwrap(); + assert_eq!(listed[0].model.as_deref(), Some("kimi-k3")); +} + #[tokio::test] async fn viewing_unloaded_task_thinking_and_history_does_not_load_context() { let harness = Harness::new().await; diff --git a/apps/maple-agent/crates/maple-agent/src/agent/thinking.rs b/apps/maple-agent/crates/maple-agent/src/agent/thinking.rs index 8374a1afa..38a797d65 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/thinking.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/thinking.rs @@ -6,32 +6,39 @@ use super::config::{load_agent_config_inner, save_agent_config_inner}; use super::{AgentRuntimeHandle, AgentThinkingState, ThinkingLevel}; impl AgentRuntimeHandle { + /// The task's session on `model_id`, or `None` while a started task is + /// on another model: it switches with its next run. async fn thinking_session( &self, session_id: &str, model_id: &str, - ) -> Result { + ) -> Result, String> { let runtime = self.runtime().await?; - if let Some(session) = runtime.loaded_session(session_id).await + let loaded = runtime.loaded_session(session_id).await; + if let Some(session) = &loaded && session.model().is_some_and(|model| model.id == model_id) { - return Ok(session); + return Ok(loaded); } let row = runtime .store .get(session_id)? .ok_or_else(|| format!("Failed to find Agent task {session_id}"))?; - let model_id = if row.message_count > 0 { - row.model.as_deref().unwrap_or(model_id) - } else { - model_id + let current = match &loaded { + Some(session) => session.model().map(|model| model.id), + None => row.model.clone(), }; + if row.message_count > 0 && current.is_some_and(|current| current != model_id) { + return Ok(None); + } let model = runtime.pi_model(model_id, None, None).await; - runtime.task_session(&row, model).await + runtime.task_session(&row, model).await.map(Some) } /// A task reads its persisted Pi level. A new-task composer previews - /// the saved default, or Pi's Medium default, clamped by its model. + /// the saved default, or Pi's Medium default, clamped by its model. A + /// task previews another model with the level a switch gives it: the + /// saved default, as Pi's switch takes its default. pub async fn thinking_state( &self, session_id: Option<&str>, @@ -40,45 +47,46 @@ impl AgentRuntimeHandle { if let Some(session_id) = session_id { self.verify_generation().await?; let runtime = self.runtime().await?; - if let Some(session) = runtime.loaded_session(session_id).await { - let model = session.model().ok_or("No model selected")?; - return Ok(AgentThinkingState { - levels: model.thinking_levels(), - current: session.thinking_level(), - }); - } - let store = self.store()?; - let facts = store - .store() - .task_facts(store.account(), session_id) - .map_err(|error| error.to_string())?; - let selected_model = if facts.message_count > 0 { - facts.model.as_deref().unwrap_or(model_id) - } else { - model_id + let (task_model, task_level) = match runtime.loaded_session(session_id).await { + Some(session) => { + let model = session.model().ok_or("No model selected")?; + if model.id == model_id { + return Ok(AgentThinkingState { + levels: model.thinking_levels(), + current: session.thinking_level(), + }); + } + (Some(model.id), Some(session.thinking_level())) + } + None => { + let store = self.store()?; + let facts = store + .store() + .task_facts(store.account(), session_id) + .map_err(|error| error.to_string())?; + let saved = facts + .thinking_level + .as_deref() + .map(|value| { + serde_json::from_value::(serde_json::Value::String( + value.to_owned(), + )) + }) + .transpose() + .map_err(|_| "Invalid saved thinking level".to_string())?; + (facts.model, saved) + } }; - let model = runtime.pi_model(selected_model, None, None).await; - let saved = facts - .thinking_level - .as_deref() - .map(|value| { - serde_json::from_value::(serde_json::Value::String( - value.to_owned(), - )) - }) - .transpose() - .map_err(|_| "Invalid saved thinking level".to_string())?; - let level = match saved { - Some(level) => level, - None => self - .load_config() - .await? - .default_thinking_level - .unwrap_or(ThinkingLevel::Medium), + let default = self.load_config().await?.default_thinking_level; + let level = if task_model.as_deref().is_none_or(|model| model == model_id) { + task_level.or(default) + } else { + default.or(task_level) }; + let model = runtime.pi_model(model_id, None, None).await; return Ok(AgentThinkingState { levels: model.thinking_levels(), - current: model.clamp_thinking_level(level), + current: model.clamp_thinking_level(level.unwrap_or(ThinkingLevel::Medium)), }); } let runtime = self.runtime().await?; @@ -104,8 +112,11 @@ impl AgentRuntimeHandle { ) -> Result { self.verify_generation().await?; self.ensure_accepting_new_work()?; - let state = if let Some(session_id) = session_id { - let session = self.thinking_session(session_id, model_id).await?; + let session = match session_id { + Some(session_id) => self.thinking_session(session_id, model_id).await?, + None => None, + }; + let state = if let Some(session) = session { session.set_thinking_level(level).await; AgentThinkingState { levels: session @@ -115,6 +126,8 @@ impl AgentRuntimeHandle { current: session.thinking_level(), } } else { + // A new task, or a task's next model, which its switch gives the + // default saved below. let runtime = self.runtime().await?; let model = runtime.pi_model(model_id, None, None).await; AgentThinkingState { diff --git a/apps/maple-agent/crates/maple-agent/src/store/tasks.rs b/apps/maple-agent/crates/maple-agent/src/store/tasks.rs index 2928dc708..d2ee44343 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/tasks.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/tasks.rs @@ -113,7 +113,8 @@ pub struct TaskPatch { pub provisional: Option, pub web_enabled: Option, pub settings: Option, - /// A selection before the first conversation. Source appends own it afterward. + /// A selection for the task's next run. Its next source append records + /// the model the session then has. pub model: Option, } @@ -330,14 +331,6 @@ impl Store { if expected.is_some_and(|expected| !same_task_metadata(expected, &row)) { return Err(StoreError::Conflict); } - if patch - .model - .as_ref() - .is_some_and(|model| row.facts.model.as_ref() != Some(model)) - && row.facts.message_count != 0 - { - return Err(StoreError::Conflict); - } let project = match &patch.project_root { Some(root) => ensure_project(tx, account, root, row.updated_ms)?, None => row.project, diff --git a/apps/maple-agent/docs/scenario-checklist.md b/apps/maple-agent/docs/scenario-checklist.md index 0731f96b2..b51a26e14 100644 --- a/apps/maple-agent/docs/scenario-checklist.md +++ b/apps/maple-agent/docs/scenario-checklist.md @@ -31,6 +31,7 @@ stdio MCP server configured in Settings > Integrations. | C4 | Attachments | Attach an image by picker and by paste, then send | The model sees the image (vision model) or a `read_image` reference (text model) | | C5 | Voice input | Dictate a message | The transcript of the dictation lands in the composer | | C6 | Per-task pickers | Switch the model, Web and an integration on the draft, then send | The created task carries the three choices | +| C7 | Switch models mid-task | In a task with history, pick another model and send; reopen the task after a restart | The reply comes from the new model with the earlier context, the transcript shows "Switched to …", and the task reopens on the new model (**automated**: `a_run_wakes_a_settled_task_and_a_started_task_switches_models`) | ## Slash commands and skills @@ -94,6 +95,7 @@ stdio MCP server configured in Settings > Integrations. | D4 | Plan meter | Settings > Usage | The plan, percent used and reset date show, nothing else | | D5 | Notifications | Finish a task and ask a question with the window in the background | Two notifications, none for permissions | | D6 | Appearance, fonts, motion, Vim | Change each setting | Each applies live | +| D8 | Compaction after a switch to a smaller window | A long task on a large-window model, then a model whose window it exceeds | The task compacts before the next reply, summarized by the previous model (log line "summarized by"), and the reply comes from the new one (**automated**: `a_switch_to_a_smaller_window_compacts_with_the_previous_model`) | | D7 | Tasks from a Goose build | Open the app on an account with tasks a Goose build saved | The app starts and runs new tasks; the old ones are not listed, since the Pi runtime reads no Goose-era data | ## Results