From b957e123705cdadba378ec302bdc60fc9dc129b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sun, 13 Sep 2026 00:15:14 -0400 Subject: [PATCH] fix(tui): repaint the header context window on a /model switch After `/model`, the header kept the previous model's window ("15K / 1.0M" on a 262K model) until the next turn or a restart. The pager caches the window inside its context snapshot (streaming `totalTokens` updates refresh it from the catalog) and nothing refreshed the cache when the session's model changed. Comparing cached numbers against catalog windows cannot tell a catalog-derived total from a shell effective limit (KIGI_DEBUG_CONTEXT_WINDOW, header upgrades), so the shell, which installs the effective window, now reports it: - kigi-shell: `handle_set_session_model` returns the installed window; `model_switch::apply` stamps it as `totalContextTokens` in the `session/set_model` response meta and carries it on the `ModelChanged` broadcast (`context_window`, serde-default: old pagers ignore it, old leaders omit it). The response meta previously serialized the actor's inner `Result` and reported success on an actor failure; that error now propagates. - kigi-tui: `Effect::SwitchModel` reads the meta; `SwitchModelComplete` carries `Ok(Option)`; the completion handler and the follower `ModelChanged` arm call `AgentView::apply_context_window`, which re-derives the snapshot (pct, free tokens) and reports whether it changed so a follower redraws on a window-only change. - Tests: the bare `/model ` dispatch path asserts the snapshot and the drawn header; follower with/without the field; same-model new window redraw; actor returns the window; serde round trip. Claude-Session: https://claude.ai/code/session_01RUqV8HVGXcAETyYbZ2tPQ4 --- .../src/agent/handlers/model_switch.rs | 16 ++- .../kigi-shell/src/extensions/notification.rs | 11 ++ .../kigi-shell/src/session/acp_session.rs | 3 + .../session/acp_session_impl/model_switch.rs | 4 +- .../model_switch_window_tests.rs | 24 +++++ .../kigi-shell/src/session/commands.rs | 4 +- .../app/acp_handler/session_notification.rs | 5 +- .../kigi-tui/src/app/acp_handler/tests/mod.rs | 3 + .../src/app/acp_handler/tests/models.rs | 76 +++++++++++-- crates/codegen/kigi-tui/src/app/actions.rs | 5 +- .../kigi-tui/src/app/agent_view/session.rs | 44 ++++++++ .../src/app/dispatch/session/lifecycle.rs | 7 +- .../kigi-tui/src/app/dispatch/tests/prompt.rs | 2 +- .../src/app/dispatch/tests/settings.rs | 2 +- .../src/app/dispatch/tests/task_result.rs | 102 ++++++++++++++++-- .../codegen/kigi-tui/src/app/effects/mod.rs | 7 +- 16 files changed, 288 insertions(+), 27 deletions(-) create mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/model_switch_window_tests.rs diff --git a/crates/codegen/kigi-shell/src/agent/handlers/model_switch.rs b/crates/codegen/kigi-shell/src/agent/handlers/model_switch.rs index 832d6cf..7284857 100644 --- a/crates/codegen/kigi-shell/src/agent/handlers/model_switch.rs +++ b/crates/codegen/kigi-shell/src/agent/handlers/model_switch.rs @@ -210,9 +210,9 @@ pub(crate) async fn apply( auto_compact_threshold_percent: new_threshold, responds_to: tx, }); - let updated_model = rx + let (updated_model, context_window) = rx .await - .map_err(|_| acp::Error::internal_error().data("failed to set session model"))?; + .map_err(|_| acp::Error::internal_error().data("failed to set session model"))??; if let Some(handle) = agent.sessions.borrow_mut().get_mut(&session_id) { handle.model_id = model_id.clone(); handle.reasoning_effort = applied_effort; @@ -224,6 +224,7 @@ pub(crate) async fn apply( &session_id, model_id.0.as_ref(), applied_effort.map(|eff| eff.to_string()), + context_window.get(), ); if agent.cfg.borrow().mode != config::AgentMode::Leader { agent.models_manager.set_current_model_id(model_id); @@ -232,9 +233,12 @@ pub(crate) async fn apply( .set_current_reasoning_effort(applied_effort); } Ok(acp::SetSessionModelResponse::new().meta( - serde_json::json!({ "model" : updated_model, }) - .as_object() - .cloned(), + serde_json::json!({ + "model": updated_model, + "totalContextTokens": context_window.get(), + }) + .as_object() + .cloned(), )) } /// Broadcast a `ModelChanged` to every client subscribed to this session so @@ -245,12 +249,14 @@ fn broadcast_model_changed( session_id: &acp::SessionId, model_id: &str, reasoning_effort: Option, + context_window: u64, ) { let notification = crate::extensions::notification::SessionNotification { session_id: session_id.clone(), update: crate::extensions::notification::SessionUpdate::ModelChanged { model_id: model_id.to_owned(), reasoning_effort, + context_window: Some(context_window), }, meta: None, }; diff --git a/crates/codegen/kigi-shell/src/extensions/notification.rs b/crates/codegen/kigi-shell/src/extensions/notification.rs index 095620c..a94bc20 100644 --- a/crates/codegen/kigi-shell/src/extensions/notification.rs +++ b/crates/codegen/kigi-shell/src/extensions/notification.rs @@ -691,6 +691,11 @@ pub enum SessionUpdate { /// does not support reasoning effort or no effort override was applied. #[serde(default, skip_serializing_if = "Option::is_none")] reasoning_effort: Option, + /// The session's effective context window after the switch, so a + /// follower repaints its header without waiting for the next turn. + /// `None` from an agent that predates the field. + #[serde(default, skip_serializing_if = "Option::is_none")] + context_window: Option, }, /// Streaming chunk of a tool call's arguments. /// @@ -1970,15 +1975,18 @@ mod tests { let with_effort = SessionUpdate::ModelChanged { model_id: "kigi-4".into(), reasoning_effort: Some("high".into()), + context_window: Some(262_144), }; let json = serde_json::to_value(&with_effort).unwrap(); assert_eq!(json["sessionUpdate"], "model_changed"); assert_eq!(json["model_id"], "kigi-4"); assert_eq!(json["reasoning_effort"], "high"); + assert_eq!(json["context_window"], 262_144); let without_effort = SessionUpdate::ModelChanged { model_id: "kigi-3".into(), reasoning_effort: None, + context_window: None, }; let json = serde_json::to_value(&without_effort).unwrap(); assert_eq!(json["sessionUpdate"], "model_changed"); @@ -1988,6 +1996,7 @@ mod tests { "reasoning_effort: None must be skipped on the wire so old pagers \ and third-party ACP clients see a smaller, no-extra-keys payload" ); + assert!(json.get("context_window").is_none()); } /// `ModelChanged` round-trips through JSON: a follower client deserializes @@ -2000,6 +2009,7 @@ mod tests { let original = SessionUpdate::ModelChanged { model_id: "kigi-4".into(), reasoning_effort: Some("medium".into()), + context_window: Some(1_000_000), }; let json_str = serde_json::to_string(&original).unwrap(); let parsed: SessionUpdate = serde_json::from_str(&json_str).unwrap(); @@ -2021,6 +2031,7 @@ mod tests { update: SessionUpdate::ModelChanged { model_id: "kigi-4".into(), reasoning_effort: None, + context_window: None, }, meta: None, }; diff --git a/crates/codegen/kigi-shell/src/session/acp_session.rs b/crates/codegen/kigi-shell/src/session/acp_session.rs index 2f74f85..f5c2b86 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session.rs @@ -1518,6 +1518,9 @@ mod load_user_prompts_tests; #[path = "acp_session_tests/memory_config_tests.rs"] mod memory_config_tests; #[cfg(test)] +#[path = "acp_session_tests/model_switch_window_tests.rs"] +mod model_switch_window_tests; +#[cfg(test)] #[path = "acp_session_tests/parallel_dispatch_tests.rs"] mod parallel_dispatch_tests; #[cfg(test)] diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs index 8f337dc..2877529 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs @@ -10,7 +10,7 @@ impl SessionActor { apply_prompt_override: bool, skip_prompt_rewrite: bool, auto_compact_threshold_percent: u8, - ) -> Result { + ) -> Result<(acp::ModelId, std::num::NonZeroU64), acp::Error> { let model_id = acp::ModelId::new(sampling_config.model.clone()); // H4: record the picker's catalog KEY as this SESSION's own selection. // `sampling_config.model` is the ambiguous bare slug; the key is what @@ -120,7 +120,7 @@ impl SessionActor { agent_name: Some(agent_name), reasoning_effort: Some(sampling_config.reasoning_effort), }); - Ok(model_id) + Ok((model_id, new_context_window)) } /// Handle [`SessionCommand::RebuildAgentForDefinition`]. /// diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/model_switch_window_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/model_switch_window_tests.rs new file mode 100644 index 0000000..34dc4d7 --- /dev/null +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/model_switch_window_tests.rs @@ -0,0 +1,24 @@ +//! `handle_set_session_model` reports the session's effective context +//! window so clients repaint their header without waiting for a turn. +use super::support::*; + +#[tokio::test(flavor = "current_thread")] +async fn set_session_model_returns_the_effective_context_window() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (actor, _rx) = build_actor().await; + let cfg = kigi_sampler::SamplerConfig { + model: "narrow".to_string(), + context_window: 128_000, + ..Default::default() + }; + let (model_id, window) = actor + .handle_set_session_model(cfg, None, false, false, true, 85) + .await + .expect("model switch"); + assert_eq!(model_id.0.as_ref(), "narrow"); + assert_eq!(window.get(), 128_000); + }) + .await; +} diff --git a/crates/codegen/kigi-shell/src/session/commands.rs b/crates/codegen/kigi-shell/src/session/commands.rs index e809a10..247b0f3 100644 --- a/crates/codegen/kigi-shell/src/session/commands.rs +++ b/crates/codegen/kigi-shell/src/session/commands.rs @@ -173,7 +173,9 @@ pub enum SessionCommand { /// `compaction.threshold_percent` (which is `Cell` so it can /// update without `&mut self`). auto_compact_threshold_percent: u8, - responds_to: oneshot::Sender>, + /// Ok: the applied model id and the session's effective context + /// window, which clients repaint from without waiting for a turn. + responds_to: oneshot::Sender>, }, /// Zero-turn harness rebuild: build a brand-new `Agent` from the /// session's `AgentRebuildSpec` and the new `AgentDefinition`, diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs b/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs index b2ca4b8..3f74b18 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs @@ -812,6 +812,7 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu XaiSessionUpdate::ModelChanged { model_id, reasoning_effort, + context_window, } => { if agent.session.model_switch_pending { tracing::debug!( @@ -847,6 +848,8 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu .session .models .set_current(new_model_id.clone(), effort); + let window_changed = + context_window.is_some_and(|total| agent.apply_context_window(total)); agent.session.user_model_preference = Some(new_model_id.clone()); let resolved_effort = agent.session.models.reasoning_effort; let actually_changed = @@ -858,7 +861,7 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu "ModelChanged broadcast applied (remote switch)" ); } - actually_changed + actually_changed || window_changed } XaiSessionUpdate::MemoryFiles { files } => { let entries = crate::views::memory_modal::build_entries(files); diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs index dbaf6e5..b8b5e04 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs @@ -1701,12 +1701,14 @@ pub(super) fn model_changed_ext( session_id: &str, model_id: &str, reasoning_effort: Option<&str>, + context_window: Option, ) -> acp::ExtNotification { let payload = SessionNotification { session_id: acp::SessionId::new(session_id), update: XaiSessionUpdate::ModelChanged { model_id: model_id.to_string(), reasoning_effort: reasoning_effort.map(String::from), + context_window, }, meta: None, }; @@ -1723,6 +1725,7 @@ pub(super) fn model_changed_ext_with_event( update: XaiSessionUpdate::ModelChanged { model_id: model_id.to_string(), reasoning_effort: None, + context_window: None, }, meta: Some(serde_json::json!({ "eventId" : event_id })), }; diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/models.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/models.rs index 839e5c8..a3e5ba0 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/models.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/models.rs @@ -110,6 +110,70 @@ ); } + fn seed_windowed_models(agent: &mut AgentView) { + for (slug, window) in [("wide", 1_000_000u64), ("narrow", 262_144)] { + let mut info = make_model_info(slug); + info.meta = serde_json::json!({ "totalContextTokens": window }) + .as_object() + .cloned(); + agent + .session + .models + .available + .insert(acp::ModelId::new(std::sync::Arc::from(slug)), info); + } + agent + .session + .models + .set_current(acp::ModelId::new(std::sync::Arc::from("wide")), None); + agent.apply_context_used(15_000, 1_000_000); + } + + /// A follower mirrors a remote switch and installs the effective window + /// the broadcast carries, like the invoking client does. + #[test] + fn model_changed_broadcast_installs_the_reported_context_window() { + let mut app = make_app_with_agent("sess-1"); + seed_windowed_models(app.agents.get_mut(&AgentId(0)).unwrap()); + + let notif = model_changed_ext("sess-1", "narrow", None, Some(262_144)); + assert!(handle_ext_notification(¬if, &mut app)); + + let snap = app.agents[&AgentId(0)].context_state.as_ref().unwrap(); + assert_eq!((snap.used, snap.total, snap.usage_pct), (15_000, 262_144, 6)); + } + + /// Same model, same effort, new effective window (the catalog was + /// reloaded and re-applied): the follower must still request a redraw. + #[test] + fn model_changed_same_model_new_window_requests_redraw() { + let mut app = make_app_with_agent("sess-1"); + seed_windowed_models(app.agents.get_mut(&AgentId(0)).unwrap()); + + let notif = model_changed_ext("sess-1", "wide", None, Some(2_000_000)); + assert!(handle_ext_notification(¬if, &mut app)); + + let snap = app.agents[&AgentId(0)].context_state.as_ref().unwrap(); + assert_eq!((snap.used, snap.total), (15_000, 2_000_000)); + + let again = model_changed_ext("sess-1", "wide", None, Some(2_000_000)); + assert!(!handle_ext_notification(&again, &mut app)); + } + + /// A broadcast from a leader that predates the field carries no window; + /// the snapshot is left for the next streaming update to refresh. + #[test] + fn model_changed_broadcast_without_window_keeps_the_snapshot() { + let mut app = make_app_with_agent("sess-1"); + seed_windowed_models(app.agents.get_mut(&AgentId(0)).unwrap()); + + let notif = model_changed_ext("sess-1", "narrow", None, None); + assert!(handle_ext_notification(¬if, &mut app)); + + let snap = app.agents[&AgentId(0)].context_state.as_ref().unwrap(); + assert_eq!((snap.used, snap.total), (15_000, 1_000_000)); + } + #[test] fn models_update_without_active_agent_uses_shell_default() { let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); @@ -232,7 +296,7 @@ let scrollback_before = agent.scrollback.len(); assert!(!agent.session.model_switch_pending); - let notif = model_changed_ext("sess-1", "kigi-4", None); + let notif = model_changed_ext("sess-1", "kigi-4", None, None); let changed = handle_ext_notification(¬if, &mut app); assert!( changed, @@ -278,7 +342,7 @@ Some(acp::ModelId::new(std::sync::Arc::from("heavy"))); assert!(!agent.session.model_switch_pending); - let notif = model_changed_ext("sess-1", "auto", None); + let notif = model_changed_ext("sess-1", "auto", None, None); let changed = handle_ext_notification(¬if, &mut app); assert!( changed, @@ -328,7 +392,7 @@ agent.session.model_switch_pending = true; let scrollback_before = agent.scrollback.len(); - let notif = model_changed_ext("sess-1", "kigi-4", None); + let notif = model_changed_ext("sess-1", "kigi-4", None, None); let changed = handle_ext_notification(¬if, &mut app); assert!( !changed, @@ -369,7 +433,7 @@ let agent = app.agents.get_mut(&AgentId(0)).unwrap(); seed_models(agent, "kigi-3", &["kigi-3", "kigi-4"]); - let notif = model_changed_ext("sess-1", "kigi-99-unknown", None); + let notif = model_changed_ext("sess-1", "kigi-99-unknown", None, None); let changed = handle_ext_notification(¬if, &mut app); assert!( !changed, @@ -400,7 +464,7 @@ let agent = app.agents.get_mut(&AgentId(0)).unwrap(); seed_models(agent, "kigi-3", &["kigi-3", "kigi-4"]); - let notif = model_changed_ext("sess-1", "kigi-4", Some("high")); + let notif = model_changed_ext("sess-1", "kigi-4", Some("high"), None); assert!(handle_ext_notification(¬if, &mut app)); let agent = app.agents.get(&AgentId(0)).unwrap(); @@ -422,7 +486,7 @@ let agent = app.agents.get_mut(&AgentId(0)).unwrap(); seed_models(agent, "kigi-3", &["kigi-3", "kigi-4"]); - let notif = model_changed_ext("sess-OTHER", "kigi-4", None); + let notif = model_changed_ext("sess-OTHER", "kigi-4", None, None); let changed = handle_ext_notification(¬if, &mut app); assert!(!changed); diff --git a/crates/codegen/kigi-tui/src/app/actions.rs b/crates/codegen/kigi-tui/src/app/actions.rs index 554990e..12f27cd 100644 --- a/crates/codegen/kigi-tui/src/app/actions.rs +++ b/crates/codegen/kigi-tui/src/app/actions.rs @@ -2112,7 +2112,10 @@ pub enum TaskResult { agent_id: AgentId, model_id: acp::ModelId, effort: Option, - result: Result<(), SwitchModelError>, + /// Ok carries the effective context window the shell reported + /// (`totalContextTokens` in the response meta; `None` from a shell + /// that predates it). + result: Result, SwitchModelError>, /// Forwarded from `Effect::SwitchModel.prev_model_id` for /// rollback on `IncompatibleAgent`. prev_model_id: Option, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/session.rs b/crates/codegen/kigi-tui/src/app/agent_view/session.rs index eb33c24..90c4ab0 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/session.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/session.rs @@ -733,6 +733,20 @@ impl AgentView { } } } + /// The shell reported the session's effective context window (model + /// switch). The snapshot caches the previous model's window, so the + /// header would show it until the next streaming update. Returns + /// whether the snapshot changed, so the caller can request a redraw. + pub fn apply_context_window(&mut self, total: u64) -> bool { + let Some((used, cached)) = self.context_state.as_ref().map(|c| (c.used, c.total)) else { + return false; + }; + if cached == total { + return false; + } + self.apply_context_used(used, total); + true + } /// Record a key event to the input flight recorder. /// /// Zero heap allocations — stores raw `Copy` types in the ring buffer. @@ -1262,6 +1276,36 @@ mod resolve_turn_activity_tests { } } #[cfg(test)] +mod apply_context_window_tests { + use super::super::test_agent_view; + fn view() -> super::AgentView { + test_agent_view(Some("s1"), std::path::PathBuf::from("/tmp")) + } + #[test] + fn reported_window_replaces_the_cached_total() { + let mut agent = view(); + agent.apply_context_used(15_000, 1_000_000); + assert!(agent.apply_context_window(262_144)); + let snap = agent.context_state.as_ref().unwrap(); + assert_eq!( + (snap.used, snap.total, snap.usage_pct, snap.free_tokens), + (15_000, 262_144, 6, 262_144 - 15_000) + ); + } + #[test] + fn same_window_reports_no_change() { + let mut agent = view(); + agent.apply_context_used(15_000, 1_000_000); + assert!(!agent.apply_context_window(1_000_000)); + } + #[test] + fn no_snapshot_stays_none() { + let mut agent = view(); + assert!(!agent.apply_context_window(262_144)); + assert!(agent.context_state.is_none()); + } +} +#[cfg(test)] mod status_window_tests { use super::super::test_agent_view; #[test] diff --git a/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs b/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs index 7e78406..a33b4dc 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs @@ -978,13 +978,13 @@ pub(in crate::app::dispatch) fn handle_switch_model_complete( agent_id: AgentId, model_id: acp::ModelId, effort: Option, - result: Result<(), SwitchModelError>, + result: Result, SwitchModelError>, prev_model_id: Option, ) -> Vec { if let Some(agent) = app.agents.get_mut(&agent_id) { agent.session.model_switch_pending = false; let mut effects = match result { - Ok(()) => { + Ok(context_window) => { agent.session.user_model_preference = Some(model_id.clone()); let display_name = agent .session @@ -996,6 +996,9 @@ pub(in crate::app::dispatch) fn handle_switch_model_complete( let prev_model = agent.session.models.current.clone(); let prev_effort = agent.session.models.reasoning_effort; agent.session.models.set_current(model_id.clone(), effort); + if let Some(total) = context_window { + agent.apply_context_window(total); + } let resolved_effort = agent.session.models.reasoning_effort; let unchanged = prev_model.as_ref() == Some(&model_id) && prev_effort == resolved_effort; diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs index 8f9c235..124489a 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs @@ -1789,7 +1789,7 @@ fn switch_model_holds_prompt_until_complete() { agent_id: id, model_id, effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs index 801ed2f..9921502 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs @@ -282,7 +282,7 @@ fn model_switch_pending_resets_correctly_across_success_and_failure() { agent_id: id, model_id: model_a, effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs index 9af1f46..4e82515 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs @@ -535,7 +535,7 @@ fn switch_model_complete_success_updates_model_and_pushes_message() { agent_id: id, model_id: model_id.clone(), effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, @@ -554,6 +554,96 @@ fn switch_model_complete_success_updates_model_and_pushes_message() { )); } +/// The bare `/model ` path: `SetDefaultModel` moves `models.current` +/// before the ACP round trip; the completion then installs the effective +/// window the shell reported, so the header repaints at once. +#[test] +fn switch_model_complete_installs_the_reported_context_window() { + let mut app = test_app_with_agent(); + let id = AgentId(0); + let model = |slug: &str, window: u64| { + let mid = acp::ModelId::new(std::sync::Arc::from(slug)); + let info = acp::ModelInfo::new(mid.clone(), slug.to_string()).meta( + serde_json::json!({ "totalContextTokens": window }) + .as_object() + .cloned(), + ); + (mid, info) + }; + let (wide, wide_info) = model("wide", 1_000_000); + let (narrow, narrow_info) = model("narrow", 262_144); + { + let agent = app.agents.get_mut(&id).unwrap(); + agent + .session + .models + .available + .insert(wide.clone(), wide_info); + agent + .session + .models + .available + .insert(narrow.clone(), narrow_info); + agent.session.models.set_current(wide.clone(), None); + agent.apply_context_used(15_000, 1_000_000); + } + + let effects = dispatch(Action::SetDefaultModel(narrow.clone()), &mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::SwitchModel { .. })) + ); + assert!(app.agents[&id].session.model_switch_pending); + assert_eq!( + app.agents[&id].context_state.as_ref().unwrap().total, + 1_000_000 + ); + + dispatch( + Action::TaskComplete(TaskResult::SwitchModelComplete { + agent_id: id, + model_id: narrow, + effort: None, + result: Ok(Some(262_144)), + prev_model_id: Some(wide), + }), + &mut app, + ); + + let agent = app.agents.get_mut(&id).unwrap(); + let snap = agent.context_state.as_ref().unwrap(); + assert_eq!( + (snap.used, snap.total, snap.usage_pct), + (15_000, 262_144, 6) + ); + let area = ratatui::layout::Rect::new(0, 0, 100, 40); + let mut buf = ratatui::buffer::Buffer::empty(area); + agent.draw( + area, + &mut buf, + &crate::actions::ActionRegistry::defaults(), + &mut crate::scrollback::render::ScratchBuffer::new(), + None, + false, + 0, + None, + &crate::app::bundle::BundleState::default(), + false, + &mut Vec::new(), + ); + let screen: String = (0..area.height) + .map(|y| { + (0..area.width) + .filter_map(|x| buf.cell((x, y)).map(|c| c.symbol().to_string())) + .collect::() + }) + .collect::>() + .join("\n"); + assert!(screen.contains("15K / 262K"), "header: {screen:?}"); + assert!(!screen.contains("1.0M"), "stale window: {screen:?}"); +} + #[test] fn switch_model_complete_skips_message_and_persist_when_unchanged() { let mut app = test_app_with_agent(); @@ -575,7 +665,7 @@ fn switch_model_complete_skips_message_and_persist_when_unchanged() { agent_id: id, model_id: model_id.clone(), effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, @@ -630,7 +720,7 @@ fn switch_model_complete_persists_resolved_effort_from_catalog_meta() { model_id: model_id.clone(), // user typed `/model Blackbox 4.7` with no effort effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, @@ -696,7 +786,7 @@ fn switch_to_non_reasoning_model_clears_persisted_effort() { agent_id: id, model_id: model_id.clone(), effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, @@ -930,7 +1020,7 @@ fn same_agent_type_switch_no_modal() { agent_id: id, model_id: model_b.clone(), effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, @@ -969,7 +1059,7 @@ fn switch_model_pending_lifecycle() { agent_id: id, model_id, effort: None, - result: Ok(()), + result: Ok(None), prev_model_id: None, }), &mut app, diff --git a/crates/codegen/kigi-tui/src/app/effects/mod.rs b/crates/codegen/kigi-tui/src/app/effects/mod.rs index 061d8c3..ae71fc0 100644 --- a/crates/codegen/kigi-tui/src/app/effects/mod.rs +++ b/crates/codegen/kigi-tui/src/app/effects/mod.rs @@ -1587,7 +1587,12 @@ pub(crate) fn execute( .meta(meta); let result = acp_send(req, &tx) .await - .map(|_| ()) + .map(|resp| { + resp.meta + .as_ref() + .and_then(|m| m.get("totalContextTokens")) + .and_then(serde_json::Value::as_u64) + }) .map_err(|e| { use kigi_shell::agent::config::ModelSwitchIncompatibleAgentError; if let Some(typed) = ModelSwitchIncompatibleAgentError::from_acp_error(