diff --git a/docs/user/tui-and-sessions.md b/docs/user/tui-and-sessions.md index f999daa5..b94bb458 100644 --- a/docs/user/tui-and-sessions.md +++ b/docs/user/tui-and-sessions.md @@ -210,15 +210,17 @@ The TUI handles `/new`, `/resume`, `/sessions`, `/close`, `/model`, `/effort`, ` The other commands listed above are available only while the session is idle. `/agents` toggles the agent roster without starting a model turn. `/new` closes the current session and starts a fresh persisted session. It clears the visible transcript but does not delete or alter the previous session, which remains resumable by its ID. Text following `/new` becomes the new session's first prompt. `/resume ` closes the current session, resumes the requested durable session, and replays its transcript; selecting the already-active ID is a no-op. `/sessions` opens a visible newest-first selector for the same workspace. Up and Down move, Enter uses the existing resume flow, `R` opens an inline rename field, and Esc cancels renaming or closes the dialog. Submit an empty rename and confirm to clear the custom name. After a save, the picker remains open on the selected session and refreshes its displayed name. `/close` closes the current session and exits the TUI. -`/model` opens the model selector on a provider list; Enter on a provider shows that provider's models. Once the session is open, providers that Kit can sign in to but that have no credentials yet are listed greyed out below the signed-in ones. Enter on one starts its terminal sign-in, the same flow as `/login `, and is available only while the session is idle; after a successful sign-in, the selector reopens on that provider's models with the same Tab setting. Typing searches only the current page, and the search clears when you move between pages. Esc returns from the model list to the provider list, and closes the selector from the provider list. `/effort` opens the advertised ACP reasoning-effort selector; `/effort default|low|medium|high` selects directly. In either dialog, Tab toggles saving the selection to `~/.kit/config.toml`, Enter selects, and Esc closes. Saving `default` removes top-level `reasoning_effort`; other values update it without replacing unrelated TOML. New sessions start from the resolved CLI/TOML default. Durable sessions remember their reasoning-effort selection, including an explicit provider `default`, and restore it on load or resume even if the process default has changed. Older transcripts without a recorded effort continue to use the process default. Forks inherit the source session’s current effort. Saving to TOML changes the default for future sessions; it is not required to retain the effort of the current session. Model selection still starts from the resolved CLI/TOML default on reload. Reasoning-effort records require a Kit version that understands transcript schema 6; older versions reject these records rather than silently discard the saved selection. +`/model` opens the model selector on a provider list; Enter on a provider shows that provider's models. Once the session is open, providers that Kit can sign in to but that have no credentials yet are listed greyed out below the signed-in ones. Enter on one starts its terminal sign-in, the same flow as `/login `, and is available only while the session is idle; after a successful sign-in, the selector reopens on that provider's models with the same Tab setting. Typing searches only the current page, and the search clears when you move between pages. Esc returns from the model list to the provider list, and closes the selector from the provider list. `/effort` opens the advertised ACP reasoning-effort selector; `/effort default|low|medium|high` selects directly. In either dialog, Tab toggles saving the selection to `~/.kit/config.toml`, Enter selects, and Esc closes. Saving `default` removes top-level `reasoning_effort`; other values update it without replacing unrelated TOML. New sessions start from the resolved CLI/TOML default. Durable sessions remember their reasoning-effort selection, including an explicit provider `default`, and restore it on load or resume even if the process default has changed. Older transcripts without a recorded effort continue to use the process default. Forks inherit the source session’s current effort. Saving to TOML changes the default for future sessions; it is not required to retain the effort of the current session. Reasoning-effort records require a Kit version that understands transcript schema 6; older versions reject these records rather than silently discard the saved selection. -Before changing models, Kit compares the latest provider-reported transcript occupancy with the target model's advertised context window. It adds a 20% tokenizer margin (`ceil(tokens × 1.20)`) and warns when that estimate is at least 80% of the target window. This is the latest request's occupancy, not accumulated session usage; cached input tokens are not counted twice. +Durable sessions also remember their model and restore it on load or resume, even if the configured default has changed. Forks inherit the source session's current model. An explicit `--model` or `--provider` on the command that resumes a session overrides the remembered model and replaces it. If Kit can show that the remembered model is unusable, because its provider has no credentials or a successful model listing omits it, the session uses the configured default for that attachment and Kit reports a warning; the remembered model stays on record and returns if it becomes available again. A failed listing proves nothing, so Kit keeps the remembered model. Older transcripts without a remembered model use the configured default. Model records require a Kit version that understands transcript schema 7; older versions reject them. + +Before changing models, Kit compares the latest provider-reported transcript occupancy with the target model's advertised context window. It adds a 20% tokenizer margin (`ceil(tokens × 1.20)`) and warns when that estimate is at least 80% of the target window. It does not warn when the target window is at least as large as the window recorded with the latest usage, because that switch cannot leave the session with less room. This is the latest request's occupancy, not accumulated session usage; cached input tokens are not counted twice. The TUI warning offers **Continue anyway**, **Compact**, and **Cancel** (the default). Continue switches without compacting. Compact runs the existing compactor with the **original model** and applies the new selection only after successful completion and durable transcript replacement. A failure or cancellation keeps the original model selected; compaction can already have changed the transcript if cancellation arrives after replacement. Cancel dismisses the warning without changing the model or transcript, and preserves the input draft. Esc cancels an in-progress switch; repeated Esc is harmless, and Ctrl+C while cancellation is pending exits if it is stuck. Stale confirmations are rejected when the session, selected model, target, or transcript changes. If the latest token count or target context window is unavailable (including fallback/custom models without catalog metadata), Kit allows an **unchecked switch** rather than inventing a limit or forcing compaction. The estimate is a warning, not a guarantee that the next provider request fits. -External ACP clients receive a confirmation-required error instead of a Kit dialog. Its extension key in error data is `kit.model_switch`; clients can resubmit the same configuration request with `_meta: {"kit.model_switch": {"token": , "action": "continue"}}`. Both ACP versions also accept `"action": "compact"` to compact transactionally with the original model. Confirmation tokens are session-local, in-memory, one-shot decisions, not persisted configuration. +External ACP clients receive a confirmation-required error (code `-31050`) instead of a Kit dialog. Its extension key in error data is `kit.model_switch`; clients can resubmit the same configuration request with `_meta: {"kit.model_switch": {"token": , "action": "continue"}}`. Both ACP versions also accept `"action": "compact"` to compact transactionally with the original model. Confirmation tokens are session-local, in-memory, one-shot decisions, not persisted configuration. The ACP server advertises `compact` for every new session. The TUI submits `/compact` unchanged like any other prompt; the runtime consumes exactly one text part beginning with the exact raw token before model dispatch and permits other client-provided context parts. Used alone, it ends after compaction. Whitespace-trimmed text following `/compact` and any other context parts are retained as the latest user message and start the next turn after compaction. Leading whitespace, near-misses such as `/compactness`, prompts containing multiple `/compact` command parts, and unknown slash commands remain ordinary prompts. Local commands win if an advertised command has the same name. @@ -251,7 +253,7 @@ Session discovery and restoration are isolated to the server's canonical workspa An arbitrary ACP load or resume never applies the server process's configured `--force` setting. The one exception is the initial resume requested by `kit tui --resume [] --force`: only the explicitly named or picker-selected session may use the stale-lock override. If another live Kit instance owns the session lock, restoration fails instead of taking over the session. A missing or invalid ID also fails normally. After the session closes and releases its lock, an ACP client can restore it again. -Before the restoration response, Kit replays the canonical transcript as ordered ACP updates for representable user text and attachments, assistant text and thoughts, and tool calls and results. Internal instructions, ambient context, notifications, and provider-specific content are not replayed to the client, but remain in the model transcript. Because compaction replaces the canonical transcript, restoring a compacted session replays its canonical summary history rather than the superseded pre-compaction items. +Before the restoration response, Kit replays the canonical transcript as ordered ACP updates for representable user text and attachments, assistant text and thoughts, and tool calls and results. Internal instructions, ambient context, notifications, and provider-specific content are not replayed to the client, but remain in the model transcript. Because compaction replaces the canonical transcript, restoring a compacted session replays its canonical summary history rather than the superseded pre-compaction items. After the restoration response, Kit sends a usage update with the latest recorded context occupancy when the context window is known: the window advertised for the selected model, or otherwise the window recorded with that usage if it belongs to that model. Changing models also refreshes this update with the new model's window when it is advertised. ### Session locks, `--resume`, and `--force` diff --git a/src/compaction.rs b/src/compaction.rs index d07f37fc..e8e4a606 100644 --- a/src/compaction.rs +++ b/src/compaction.rs @@ -1275,7 +1275,7 @@ fn compaction_reason(transcript: &[Item]) -> Option { configured_reason(transcript, &CompactionConfig::default()) } -fn context_window(transcript: &[Item]) -> Option { +pub(crate) fn context_window(transcript: &[Item]) -> Option { transcript .iter() .rev() diff --git a/src/main.rs b/src/main.rs index ce4f0652..e5162423 100644 --- a/src/main.rs +++ b/src/main.rs @@ -848,6 +848,13 @@ impl Command { .help("Do not serve ACP on stdio. Requires remote ACP over HTTP") .requires("remote_acp"), ); + let command = command.arg( + clap::Arg::new("prefer_session_model") + .long("prefer-session-model") + .action(clap::ArgAction::SetTrue) + .help("Treat --model and --provider as defaults for resumed sessions") + .hide(true), + ); let command = command.arg( clap::Arg::new("stdio_protocol_version") .long("stdio-protocol-version") @@ -1190,6 +1197,7 @@ impl Command { no_a2a: required_arg(matches, "no_a2a")?, no_stdio: required_arg(matches, "no_stdio")?, stdio_protocol_version: required_arg(matches, "stdio_protocol_version")?, + prefer_session_model: required_arg(matches, "prefer_session_model")?, server_credential_file: optional_arg(matches, "server_credential_file")?, mcp: McpArgs::from_matches(matches)?, terminal_auth_login: optional_arg(matches, "terminal_auth_login")?, @@ -1770,6 +1778,8 @@ enum Command { no_stdio: bool, /// ACP wire version for stdio (defaults to v1 for compatibility). stdio_protocol_version: AcpProtocolVersion, + /// Internal: `--model` and `--provider` are only defaults for resumed sessions. + prefer_session_model: bool, /// Require this file's bearer token on every HTTP request. server_credential_file: Option, @@ -2317,6 +2327,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { no_a2a, no_stdio, stdio_protocol_version, + prefer_session_model, server_credential_file, mcp, terminal_auth_login: _, @@ -2333,6 +2344,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { return Ok(()); } let root = config.root(root); + let model_override = (model.is_some() || provider.is_some()) && !prefer_session_model; let model = config.model(model); let provider = config.provider(provider); let reasoning_effort = config.reasoning_effort(reasoning_effort); @@ -2365,6 +2377,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { runtime, config.inference_providers.clone(), )?; + let runtime = kit::Runtime::with_model_override(runtime, model_override)?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_skill_poll_interval( runtime, @@ -2447,6 +2460,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { return Ok(()); } let root = config.root(root); + let model_override = model.is_some() || provider.is_some(); let model = config.model(model); let provider = config.provider(provider); let reasoning_effort = config.reasoning_effort(reasoning_effort); @@ -2478,6 +2492,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { runtime, config.inference_providers.clone(), )?; + let runtime = kit::Runtime::with_model_override(runtime, model_override)?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_skill_poll_interval( runtime, @@ -2537,6 +2552,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { } => { let (config, telemetry_settings, _telemetry, openrouter_api_key) = initialize.await?; let root = config.root(root); + let model_override = model.is_some() || provider.is_some(); let model = config.model(model); let provider = config.provider(provider); let reasoning_effort = config.reasoning_effort(reasoning_effort); @@ -2562,6 +2578,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { runtime, config.inference_providers.clone(), )?; + let runtime = kit::Runtime::with_model_override(runtime, model_override)?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_skill_poll_interval( runtime, @@ -2658,6 +2675,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { Some(Some(id)) => Some(id), None => None, }; + let model_override = model.is_some() || provider.is_some(); let model = config.model(model); let provider = config.provider(provider); let reasoning_effort = config.reasoning_effort(reasoning_effort); @@ -2696,6 +2714,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { resume.as_deref(), force, voice_enabled, + model_override, &mut stop, ) .await? @@ -3324,6 +3343,7 @@ mod tests { let Command::Serve { a2a, stdio_protocol_version, + prefer_session_model, remote_acp, no_a2a, no_stdio, @@ -3341,6 +3361,7 @@ mod tests { super::AcpProtocolVersion::V1 )); assert!(!remote_acp && !no_a2a && !no_stdio && !resume && !force); + assert!(!prefer_session_model); assert!(!mcp.no_configured_mcp_config && !mcp.legacy_mcp_config); assert!(mcp.credentials.credential_store.is_none()); let cli = Cli::try_parse_from(["kit", "acp"]).unwrap(); diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 60e04c01..25cbdf2a 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -10,6 +10,7 @@ use std::{ time::Duration, }; +use agent_client_protocol::schema::v1::UsageUpdate; use agent_client_protocol::{Client, ConnectTo, ConnectionTo, Handled}; use agenthooksprotocol::SessionEndInputOutcome; use agentkit_acp::{ @@ -1966,6 +1967,13 @@ impl Server { .map_err(|error| record_acp_runtime_failure(&session_id, "reasoning_effort", error))?; let catalog = driver.adapter.model_catalog(¤t).await; let config_options = config_options(¤t, reasoning_effort, &catalog); + let model_fallback = driver.model_fallback.clone(); + let restored_context = usage::restored_context( + &driver.canonical_transcript, + &catalog, + ¤t, + driver.unattributed_usage_applies, + ); let background_jobs = driver.background_jobs.clone(); let tasks = driver.tasks.clone(); let structured_completion = driver.structured_completion; @@ -2015,6 +2023,20 @@ impl Server { if activated.await.is_ok() { let runtime = Arc::clone(&actor.runtime); let session_id = actor.session_id.clone(); + // Clients discard non-transcript updates during replay, so these + // follow the response rather than joining the replay stream. + if let Some((used, size)) = restored_context { + let _ = connection.send_notification(SessionNotification::new( + session_id.clone(), + SessionUpdate::UsageUpdate(UsageUpdate::new(used, size)), + )); + } + if let Some(notice) = model_fallback.filter(|_| notices_supported) { + let _ = connection.send_notification(SessionNotification::new( + session_id.clone(), + SessionUpdate::Notice(Notice::new(NoticeSeverity::Warning, notice)), + )); + } let actor = std::pin::pin!(subagents::with_updates( session_actor(actor), subagent_updates, @@ -2432,6 +2454,7 @@ async fn session_actor(actor: SessionActor) { cancellation_generation, }) => { let span = crate::telemetry::incoming::operation(request.meta.as_ref()); + let model_change = request.config_id.to_string() == MODEL_CONFIG_ID; let result = async { let cancellation = integration .cancellation_handle(&session_id) @@ -2439,7 +2462,7 @@ async fn session_actor(actor: SessionActor) { if cancellation.is_cancelled_since(cancellation_generation) { return Err(model_switch::error("model change cancelled")); } - if request.config_id.to_string() == MODEL_CONFIG_ID { + if model_change { let target = request.value.as_value_id().ok_or_else(|| { model_switch::error("selection requires an id value") })?; @@ -2513,6 +2536,22 @@ async fn session_actor(actor: SessionActor) { { tracing::debug!(%error, "failed to queue ACP config options update"); } + if result.is_ok() + && model_change + && let Ok(selection) = adapter.selection() + && let Some((used, size)) = usage::restored_context( + &driver.snapshot().transcript, + &catalog, + &selection, + false, + ) + && let Err(error) = client.notify_session(SessionNotification::new( + session_id.clone(), + SessionUpdate::UsageUpdate(UsageUpdate::new(used, size)), + )) + { + tracing::debug!(%error, "failed to queue ACP usage update"); + } let _ = reply.send(result); } Some(Command::Fork { @@ -7649,6 +7688,79 @@ pub(super) mod tests { assert_eq!(options[1]["currentValue"], "default"); } + #[tokio::test] + async fn load_reports_recorded_context_usage() { + let root = tempfile::tempdir().unwrap(); + let session_id = crate::session::new_id(); + let mut usage = agentkit_core::Usage::new(agentkit_core::TokenUsage::new(90, 10)); + usage.metadata = agentkit_core::MetadataMap::from([("context_window".into(), 1000.into())]); + drop( + crate::session::open( + root.path(), + &session_id, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "remember me"), + Item::text(ItemKind::Assistant, "remembered").with_usage(usage), + ], + ) + .unwrap(), + ); + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_credentials_and_effort( + root.path(), + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + None, + ) + .unwrap(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(serve_transport(runtime, agent_transport)); + let (updates_tx, mut updates_rx) = mpsc::unbounded_channel(); + let workspace = root.path().to_path_buf(); + agent_client_protocol::Client + .builder() + .on_receive_notification( + async move |notification: SessionNotification, _cx| { + let _ = updates_tx.send(notification.update); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, async move |connection| { + connection + .send_request(InitializeRequest::new(ProtocolVersion::V1)) + .block_task() + .await?; + connection + .send_request(LoadSessionRequest::new(session_id.clone(), workspace)) + .block_task() + .await?; + loop { + let update = timeout(Duration::from_secs(10), updates_rx.recv()) + .await + .expect("usage update timed out") + .unwrap(); + if let SessionUpdate::UsageUpdate(usage) = update { + assert_eq!((usage.used, usage.size), (100, 1000)); + break; + } + } + connection + .send_request(CloseSessionRequest::new(session_id)) + .block_task() + .await?; + Ok(()) + }) + .await + .unwrap(); + server.abort(); + } + #[tokio::test] async fn close_then_load_replays_before_response_and_returns_config_options() { let root = tempfile::tempdir().unwrap(); diff --git a/src/protocols/acp/model_switch.rs b/src/protocols/acp/model_switch.rs index a882a082..ae481d44 100644 --- a/src/protocols/acp/model_switch.rs +++ b/src/protocols/acp/model_switch.rs @@ -11,6 +11,8 @@ use serde_json::{Map, Value}; use crate::provider::{ModelGroup, ModelSelection}; pub(crate) const META: &str = "kit.model_switch"; +// JSON-RPC and ACP reserve -32768..=-32000 and -32899..=-32800. +const CONFIRMATION_REQUIRED_CODE: i32 = -31050; static NEXT_TOKEN: AtomicU64 = AtomicU64::new(1); #[derive(Clone, Debug, Deserialize, Serialize)] @@ -100,7 +102,11 @@ impl Guard { let Some((tokens, window)) = guarded_tokens(&transcript).zip(window) else { return Ok(Decision::Continue); }; - if current == &target || tokens * 100 < u128::from(window) * 80 { + // A target at least as large as the window that produced the latest usage + // cannot leave the session worse off, even when the prior model is unknown. + let fits_recorded_window = crate::compaction::context_window(&transcript) + .is_some_and(|recorded| window >= recorded); + if current == &target || fits_recorded_window || tokens * 100 < u128::from(window) * 80 { return Ok(Decision::Continue); } let token = NEXT_TOKEN.fetch_add(1, Ordering::Relaxed); @@ -124,8 +130,11 @@ impl Guard { ("target_window".into(), Value::from(warning.target_window)), ])), )])); - Err(error("target model context is at least 80% occupied after a 20% tokenizer margin; continue explicitly, compact with the current model, or cancel") - .data(data)) + Err(Error::new( + CONFIRMATION_REQUIRED_CODE, + "target model context is at least 80% occupied after a 20% tokenizer margin; continue explicitly, compact with the current model, or cancel", + ) + .data(data)) } } @@ -218,6 +227,30 @@ mod tests { assert!(check(&mut guard, vec![measured(u64::MAX)], Some(u64::MAX), None).is_err()); } + #[test] + fn model_switch_to_a_window_at_least_as_large_as_recorded_needs_no_confirmation() { + let recorded = |window: u64| { + let mut usage = Usage::new(TokenUsage::new(100, 0)); + usage.metadata = MetadataMap::from_iter([("context_window".into(), json!(window))]); + vec![Item::text(ItemKind::Assistant, "response").with_usage(usage)] + }; + for window in [150, 149] { + assert_eq!( + check(&mut Guard::default(), recorded(window), Some(150), None).unwrap(), + Decision::Continue + ); + } + assert!(check(&mut Guard::default(), recorded(151), Some(150), None).is_err()); + } + + #[test] + fn model_switch_warning_explains_itself_with_a_distinct_code() { + let error = check(&mut Guard::default(), vec![measured(100)], Some(150), None).unwrap_err(); + assert_eq!(i32::from(error.code), CONFIRMATION_REQUIRED_CODE); + assert!(error.message.contains("at least 80% occupied")); + assert!(error.data.unwrap().get(META).is_some()); + } + #[test] fn model_switch_latest_occupancy_not_lifetime_or_double_cached_tokens() { let last = Item::text(ItemKind::Assistant, "response").with_usage(Usage::new( diff --git a/src/protocols/acp/usage.rs b/src/protocols/acp/usage.rs index e883ff77..ee61f24f 100644 --- a/src/protocols/acp/usage.rs +++ b/src/protocols/acp/usage.rs @@ -128,6 +128,39 @@ impl SessionUsage { } } +/// Context occupancy for a session whose window is known without a new response. +/// A recorded window describes only the model that produced the latest usage. +pub(super) fn restored_context( + transcript: &[agentkit_core::Item], + catalog: &[crate::provider::ModelGroup], + selection: &crate::provider::ModelSelection, + unattributed_usage_applies: bool, +) -> Option<(u64, u64)> { + let used = crate::compaction::latest_context_tokens(transcript)?; + let advertised = catalog + .iter() + .find(|group| group.provider == selection.provider) + .and_then(|group| group.context_windows.get(&selection.model)) + .copied() + .filter(|window| *window > 0); + let recorded = || { + let producer = transcript + .iter() + .rev() + .find_map(|item| item.usage.as_ref())? + .metadata + .get(crate::provider::USAGE_MODEL_METADATA_KEY); + let applies = match producer { + Some(producer) => producer.as_str() == Some(selection.id().as_str()), + None => unattributed_usage_applies, + }; + applies + .then(|| crate::compaction::context_window(transcript)) + .flatten() + }; + Some((used, advertised.or_else(recorded)?)) +} + #[derive(Clone)] pub(crate) enum Accounting { Main(SessionActivity), @@ -385,6 +418,77 @@ mod tests { }) } + #[test] + fn restored_context_prefers_the_advertised_window_of_the_selected_model() { + let mut usage = Usage::new(agentkit_core::TokenUsage::new(90, 10)); + usage.metadata = MetadataMap::from([("context_window".into(), 1000.into())]); + let transcript = vec![ + agentkit_core::Item::text(agentkit_core::ItemKind::Assistant, "done").with_usage(usage), + ]; + let selection = + crate::provider::ModelSelection::new(crate::ProviderKind::OpenRouter, "selected/model"); + let catalog = |windows: Vec<(String, u64)>| { + vec![crate::provider::ModelGroup { + provider: crate::ProviderKind::OpenRouter, + models: vec!["selected/model".into()], + context_windows: windows.into_iter().collect(), + }] + }; + assert_eq!( + restored_context(&transcript, &catalog(vec![]), &selection, true), + Some((100, 1000)) + ); + assert_eq!( + restored_context( + &transcript, + &catalog(vec![("selected/model".into(), 2000)]), + &selection, + true + ), + Some((100, 2000)) + ); + assert_eq!( + restored_context(&transcript, &catalog(vec![]), &selection, false), + None + ); + let stamped = |producer: &str| { + let mut usage = Usage::new(agentkit_core::TokenUsage::new(90, 10)); + usage.metadata = MetadataMap::from([ + ("context_window".into(), 1000.into()), + ( + crate::provider::USAGE_MODEL_METADATA_KEY.into(), + producer.into(), + ), + ]); + vec![ + agentkit_core::Item::text(agentkit_core::ItemKind::Assistant, "done") + .with_usage(usage), + ] + }; + assert_eq!( + restored_context( + &stamped(&selection.id()), + &catalog(vec![]), + &selection, + false + ), + Some((100, 1000)) + ); + assert_eq!( + restored_context( + &stamped("openrouter:other/model"), + &catalog(vec![]), + &selection, + true + ), + None + ); + assert_eq!( + restored_context(&[], &catalog(vec![]), &selection, true), + None + ); + } + fn start(activity: &SessionActivity) { activity.observe(&AgentEvent::TurnStarted { session_id: SessionId::new("usage"), diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index f7560efc..fc37ddff 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -1292,6 +1292,13 @@ impl Server { .map_err(AcpRuntimeError::Loop)?; let catalog = driver.adapter.model_catalog(¤t).await; let config_options = v2_config_options(¤t, reasoning, &catalog); + let model_fallback = driver.model_fallback.clone(); + let restored_context = super::usage::restored_context( + &driver.canonical_transcript, + &catalog, + ¤t, + driver.unattributed_usage_applies, + ); let canonical_transcript = driver.canonical_transcript; // Validate before spawning or publishing the attachment. BindingGuard and // the uncommitted claim release ownership on rejection. @@ -1348,6 +1355,21 @@ impl Server { let sink = actor.sink.clone(); let runtime = Arc::clone(&actor.runtime); let session_id = actor.session_id.clone(); + if let Some((used, size)) = restored_context { + let _ = sink.update(wire::UpdateSessionNotification::new( + session_id.clone(), + wire::SessionUpdate::UsageUpdate(wire::UsageUpdate::new(used, size)), + )); + } + if let Some(notice) = model_fallback { + let _ = sink.update(wire::UpdateSessionNotification::new( + session_id.clone(), + wire::SessionUpdate::Notice(wire::Notice::new( + wire::NoticeSeverity::Warning, + notice, + )), + )); + } let actor = std::pin::pin!(super::subagents::with_updates( session_actor(actor), Some(subagent_updates), @@ -1899,6 +1921,7 @@ async fn session_actor(actor: SessionActor) cancellation_generation, }) => { let span = crate::telemetry::incoming::operation(request.meta.as_ref()); + let model_change = request.config_id.to_string() == super::MODEL_CONFIG_ID; let result = async { if handle .cancellation_handle() @@ -1976,6 +1999,22 @@ async fn session_actor(actor: SessionActor) { tracing::debug!(%error, "failed to queue ACP v2 config options update"); } + if result.is_ok() + && model_change + && let Ok(selection) = adapter.selection() + && let Some((used, size)) = super::usage::restored_context( + &driver.snapshot().transcript, + &catalog, + &selection, + false, + ) + && let Err(error) = sink.update(wire::UpdateSessionNotification::new( + session_id.clone(), + wire::SessionUpdate::UsageUpdate(wire::UsageUpdate::new(used, size)), + )) + { + tracing::debug!(%error, "failed to queue ACP v2 usage update"); + } let _ = reply.send(result); } Some(Command::Fork { @@ -7686,6 +7725,78 @@ mod tests { .expect("snapshot client failed"); } + #[tokio::test] + async fn v2_resume_reports_recorded_context_usage() { + let root = tempfile::tempdir().unwrap(); + let workspace = std::fs::canonicalize(root.path()).unwrap(); + let source = crate::session::new_id(); + let mut usage = agentkit_core::Usage::new(agentkit_core::TokenUsage::new(90, 10)); + usage.metadata = MetadataMap::from([("context_window".into(), json!(1000))]); + drop( + crate::session::open( + &workspace, + &source, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::Assistant, "remembered").with_usage(usage), + ], + ) + .unwrap(), + ); + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_and_credentials( + &workspace, + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + ) + .unwrap(); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime, SessionRegistry::new()); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let (updates_tx, mut updates_rx) = mpsc::unbounded_channel(); + agent_client_protocol::Client + .v2() + .on_receive_notification( + async move |notification: wire::UpdateSessionNotification, _cx| { + let _ = updates_tx.send(notification.update); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("usage-test", "0"), + )) + .block_task() + .await?; + cx.send_request(wire::ResumeSessionRequest::new(source.clone(), workspace)) + .block_task() + .await?; + loop { + let update = timeout(Duration::from_secs(10), updates_rx.recv()) + .await + .expect("usage update timed out") + .unwrap(); + if let wire::SessionUpdate::UsageUpdate(usage) = update { + assert_eq!((usage.used, usage.size), (100, 1000)); + break; + } + } + cx.send_request(wire::CloseSessionRequest::new(wire::SessionId::new(source))) + .block_task() + .await?; + Ok(()) + }) + .await + .unwrap(); + server.abort(); + } + #[tokio::test] async fn v2_fork_preserves_history_and_advertises_capability() { let root = tempfile::tempdir().unwrap(); diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index 914f92b5..1a6d5f8c 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -150,6 +150,8 @@ pub struct ModelGroup { pub context_windows: std::collections::HashMap, } +pub(crate) const USAGE_MODEL_METADATA_KEY: &str = "kit.model"; + #[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum ReasoningEffort { @@ -279,6 +281,7 @@ impl SelectableAdapter { ) -> Result { if persist_initial { observer.set_reasoning_effort(self.reasoning_effort()?)?; + observer.set_model(&self.selection()?.id())?; } self.session_observer = Some(observer); Ok(self) @@ -358,6 +361,7 @@ impl SelectableAdapter { .map(|catalog| DiscoveredModels { models: catalog.visible_models().to_vec(), context_windows: catalog.context_windows().clone(), + hidden_models: catalog.known_models().clone(), }) .unwrap_or_default() } @@ -406,6 +410,55 @@ impl SelectableAdapter { providers } + /// False only on evidence: no credentials, or a successful listing that omits + /// the model. A failed lookup proves nothing, so the selection is kept. + pub(crate) async fn can_restore(&self, selection: &ModelSelection) -> bool { + if self.inference.get(selection.provider).is_some() { + return true; + } + if !self.catalog_providers().contains(&selection.provider) { + return false; + } + match selection.provider { + ProviderKind::OpenAiSubscription => { + listed(&self.discovered_openai_models().await, &selection.model) + } + ProviderKind::OpenRouter => { + let url = match std::env::var_os("OPENROUTER_BASE_URL") { + None => catalog_models_url(None), + Some(url) => url.to_str().and_then(|url| catalog_models_url(Some(url))), + }; + match url { + Some(url) => fetch_model_ids(&url) + .await + .map_or(true, |models| listed(&models, &selection.model)), + None => true, + } + } + ProviderKind::Speakeasy => fetch_model_ids(OPENROUTER_MODELS_URL) + .await + .map_or(true, |models| listed(&models, &selection.model)), + ProviderKind::Anthropic => { + let models = discover_anthropic(&self.credential_storage).await; + if models.is_empty() || models.iter().any(|model| model.id == selection.model) { + return true; + } + // Listings carry canonical ids, but aliases are callable too. + let Ok(Some(credentials)) = + super::anthropic_auth::credentials(&self.credential_storage) + else { + return true; + }; + !super::anthropic_catalog::confirmed_missing( + super::anthropic_catalog::MODELS_URL, + credentials.api_key(), + &selection.model, + ) + .await + } + } + } + pub fn select(&self, selection: ModelSelection) -> Result<(), String> { if !valid_model_id(&selection.model) { return Err("model name is outside canonical bounds".into()); @@ -424,6 +477,10 @@ impl SelectableAdapter { .selection .lock() .map_err(|_| "session selection lock is poisoned")?; + // Same lock order and failure semantics as `select_reasoning_effort`. + if let Some(observer) = &self.session_observer { + observer.set_model(&selection.id())?; + } current.model = selection; current.revision = current.revision.wrapping_add(1); Ok(()) @@ -537,7 +594,7 @@ fn expose_background_call_ids(request: &mut TurnRequest) { #[async_trait] impl ModelSession for SelectableSession { - type Turn = KitTurn; + type Turn = SelectableTurn; fn set_retry_observer(&mut self, observer: Option>) { self.inner.set_retry_observer(observer.clone()); @@ -570,7 +627,11 @@ impl ModelSession for SelectableSession { self.replace_active(selected, replacement)?; } expose_background_call_ids(&mut request); - self.inner.begin_turn(request, cancellation).await + let inner = self.inner.begin_turn(request, cancellation).await?; + Ok(SelectableTurn { + inner, + model: self.active.model.id(), + }) } fn model_name(&self) -> Option<&str> { @@ -582,6 +643,36 @@ impl ModelSession for SelectableSession { } } +/// Records which model produced each usage, since the selection can later change. +pub struct SelectableTurn { + inner: KitTurn, + model: String, +} + +#[async_trait] +impl ModelTurn for SelectableTurn { + fn on_cancelled(&mut self) { + self.inner.on_cancelled(); + } + + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + let mut event = self.inner.next_event(cancellation).await?; + stamp_usage_model(&mut event, &self.model); + Ok(event) + } +} + +fn stamp_usage_model(event: &mut Option, model: &str) { + stamp_event_usage(event, |usage| { + usage + .metadata + .insert(USAGE_MODEL_METADATA_KEY.into(), model.into()); + }); +} + impl SelectableSession { fn replace_active( &mut self, @@ -1535,14 +1626,17 @@ pub(super) fn stamp_context_window( context_window: u64, provider_key: &'static str, ) { - let stamp = |usage: &mut Usage| { + stamp_event_usage(event, |usage| { usage .metadata .insert("context_window".into(), context_window.into()); usage .metadata .insert(provider_key.into(), context_window.into()); - }; + }); +} + +fn stamp_event_usage(event: &mut Option, stamp: impl Fn(&mut Usage)) { match event { Some(ModelTurnEvent::Usage(usage)) => stamp(usage), Some(ModelTurnEvent::Finished(result)) => { @@ -1938,6 +2032,13 @@ fn openrouter_fallback() -> Vec { struct DiscoveredModels { models: Vec, context_windows: std::collections::HashMap, + hidden_models: std::collections::HashSet, +} + +fn listed(discovered: &DiscoveredModels, model: &str) -> bool { + discovered.models.is_empty() + || discovered.models.iter().any(|listed| listed == model) + || discovered.hidden_models.contains(model) } fn fallback_catalog() -> DiscoveredModels { @@ -2006,6 +2107,7 @@ fn parse_discovered_models(value: &Value) -> Result { Ok(DiscoveredModels { models, context_windows, + ..Default::default() }) } @@ -2267,6 +2369,54 @@ mod tests { assert!(error.contains("cannot be empty"), "{error}"); } + #[test] + fn usage_records_the_model_that_produced_it() { + use agentkit_core::{FinishReason, MetadataMap, TokenUsage, Usage}; + use agentkit_loop::{ModelTurnEvent, ModelTurnResult}; + let stamped = |usage: &Usage| usage.metadata.get(super::USAGE_MODEL_METADATA_KEY).cloned(); + let usage = Usage::new(TokenUsage::new(1, 1)); + let mut event = Some(ModelTurnEvent::Usage(usage.clone())); + super::stamp_usage_model(&mut event, "openrouter:a"); + let Some(ModelTurnEvent::Usage(result)) = &event else { + panic!("usage event"); + }; + assert_eq!(stamped(result), Some("openrouter:a".into())); + let mut event = Some(ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "done").with_usage(usage.clone())], + usage: Some(usage), + metadata: MetadataMap::new(), + model: None, + response_id: None, + })); + super::stamp_usage_model(&mut event, "openrouter:b"); + let Some(ModelTurnEvent::Finished(result)) = &event else { + panic!("finished event"); + }; + assert_eq!( + stamped(result.usage.as_ref().unwrap()), + Some("openrouter:b".into()) + ); + assert_eq!( + stamped(result.output_items[0].usage.as_ref().unwrap()), + Some("openrouter:b".into()) + ); + } + + #[test] + fn only_a_successful_listing_can_prove_a_model_unavailable() { + use super::{DiscoveredModels, listed}; + assert!(listed(&DiscoveredModels::default(), "any")); + let discovered = DiscoveredModels { + models: vec!["visible".into()], + hidden_models: ["hidden".to_string()].into_iter().collect(), + ..Default::default() + }; + assert!(listed(&discovered, "visible")); + assert!(listed(&discovered, "hidden")); + assert!(!listed(&discovered, "retired")); + } + #[test] fn selectable_adapter_defers_openrouter_credentials_until_session_start() { let directory = tempfile::tempdir().unwrap(); diff --git a/src/provider/anthropic_catalog.rs b/src/provider/anthropic_catalog.rs index 16344dcd..24c377e3 100644 --- a/src/provider/anthropic_catalog.rs +++ b/src/provider/anthropic_catalog.rs @@ -62,6 +62,13 @@ pub(super) async fn retrieve(url: &str, key: &str, model: &str) -> Result bool { + retrieve(url, key, model) + .await + .is_err_and(|error| error.contains("HTTP 404")) +} + async fn read_model_response( client: &reqwest::Client, url: reqwest::Url, @@ -206,6 +213,20 @@ mod tests { assert!(!requests[0].contains("limit=")); } + #[tokio::test] + async fn only_a_not_found_response_confirms_a_model_is_missing() { + for (status, body, missing) in [ + (404, "{}".to_owned(), true), + (401, "{}".to_owned(), false), + (500, "{}".to_owned(), false), + (200, json!({"id":"canonical-id"}).to_string(), false), + ] { + let (url, handle) = server(vec![(status, body)]); + assert_eq!(confirmed_missing(&url, "test", "alias").await, missing); + let _ = handle.join(); + } + } + #[tokio::test] async fn retrieve_rejects_redirects_errors_and_oversized_metadata() { for (status, body) in [ diff --git a/src/provider/chatgpt.rs b/src/provider/chatgpt.rs index 65e277a8..06ce4ab7 100644 --- a/src/provider/chatgpt.rs +++ b/src/provider/chatgpt.rs @@ -116,6 +116,7 @@ impl SubscriptionModelCatalogCache { pub(crate) struct SubscriptionModelCatalog { context_windows: HashMap, visible_models: Vec, + known_models: HashSet, } impl SubscriptionModelCatalog { @@ -126,6 +127,10 @@ impl SubscriptionModelCatalog { pub(crate) fn visible_models(&self) -> &[String] { &self.visible_models } + + pub(crate) fn known_models(&self) -> &HashSet { + &self.known_models + } } fn subscription_resilience() -> ResilienceConfig { @@ -1141,6 +1146,7 @@ fn parse_model_catalog(value: &Value) -> Result Result Result, + /// Set when a resumed session's recorded model could not be restored. + pub model_fallback: Option, + /// Usage without a producing model predates model records, so it can only be + /// attributed to the default when the session has no record or override. + pub unattributed_usage_applies: bool, +} + +struct RestoredModel { + selection: ModelSelection, + fallback: Option, + unattributed_usage_applies: bool, } struct McpInstallSources { @@ -761,6 +772,7 @@ pub struct Runtime { adapter: SelectableAdapter, provider: ProviderKind, model: String, + model_override: bool, reasoning_effort: Option, credential_storage: crate::credentials::CredentialStorage, openrouter_api_key: Option, @@ -884,6 +896,7 @@ impl Runtime { adapter, provider, model, + model_override: false, reasoning_effort, credential_storage, openrouter_api_key, @@ -1036,6 +1049,52 @@ impl Runtime { } /// Configures compaction before publishing the runtime to any sessions. + /// Makes the configured model replace a resumed session's recorded model. + pub fn with_model_override( + runtime: Arc, + model_override: bool, + ) -> Result, String> { + let mut runtime = Arc::try_unwrap(runtime) + .map_err(|_| "could not configure model after runtime was shared".to_string())?; + runtime.model_override = model_override; + Ok(Arc::new(runtime)) + } + + async fn restored_model( + &self, + opened: &crate::session::OpenSession, + resume: bool, + ) -> Result { + let configured = ModelSelection::new(self.provider, self.model.clone()); + let recorded = opened.model.as_deref(); + let (selection, fallback) = match recorded { + Some(recorded) if !self.model_override && recorded != configured.id() => { + match ModelSelection::from_id(recorded) { + Ok(selection) if self.adapter.can_restore(&selection).await => { + (selection, None) + } + _ => { + let notice = format!( + "session model {recorded} is unavailable; using {}", + configured.id() + ); + tracing::warn!("{notice}"); + (configured, Some(notice)) + } + } + } + _ => (configured, None), + }; + if resume && self.model_override { + opened.observer.set_model(&selection.id())?; + } + Ok(RestoredModel { + selection, + fallback, + unattributed_usage_applies: recorded.is_none() && !self.model_override, + }) + } + pub fn with_compaction( runtime: Arc, policy: crate::compaction::policy::CompactionConfig, @@ -1552,6 +1611,7 @@ impl Runtime { adapter: self.adapter.clone(), provider: self.provider, model: self.model.clone(), + model_override: self.model_override, reasoning_effort: self.reasoning_effort, credential_storage: self.credential_storage.clone(), openrouter_api_key: self.openrouter_api_key.clone(), @@ -1914,7 +1974,10 @@ impl Runtime { error, ) })?; - let selection = self.adapter.selection()?; + let selection = self + .restored_model(&opened, request.resume) + .await? + .selection; let adapter = SelectableAdapter::new_with_credentials_effort_and_openrouter_key( selection.provider, selection.model, @@ -2334,12 +2397,22 @@ impl Runtime { } // Every ACP route owns its model selection. Changing one session // cannot redirect another session served by the same runtime. - let (selection, reasoning_effort) = selected.unwrap_or_else(|| { - ( - ModelSelection::new(self.provider, self.model.clone()), - opened.reasoning_effort.unwrap_or(self.reasoning_effort), - ) - }); + let (selection, reasoning_effort, model_fallback, unattributed_usage_applies) = + match selected { + Some((selection, reasoning_effort)) => (selection, reasoning_effort, None, false), + None => { + let restored = self + .restored_model(&opened, request.resume) + .await + .map_err(AcpRuntimeError::Loop)?; + ( + restored.selection, + opened.reasoning_effort.unwrap_or(self.reasoning_effort), + restored.fallback, + restored.unattributed_usage_applies, + ) + } + }; let adapter = SelectableAdapter::new_with_credentials_effort_and_openrouter_key( selection.provider, selection.model, @@ -2444,6 +2517,8 @@ impl Runtime { structured_completion: self.base_depth > 0, adapter, canonical_transcript, + model_fallback, + unattributed_usage_applies, }; Ok(driver) } diff --git a/src/runtime/tests.rs b/src/runtime/tests.rs index ee323cf5..a97b0478 100644 --- a/src/runtime/tests.rs +++ b/src/runtime/tests.rs @@ -228,6 +228,209 @@ fn resolved_reasoning_effort_reaches_root_adapter_and_kit_children() { ); } +mod model_persistence { + use super::*; + use crate::runtime::{AcpDriver, AcpDriverContext, SessionClaim}; + + #[derive(Clone)] + struct Observer; + + impl agentkit_loop::LoopObserver for Observer { + fn handle_event(&self, _event: agentkit_loop::ObservedEvent) {} + } + + fn runtime(root: &std::path::Path, model: &str, model_override: bool) -> Arc { + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_credentials_and_effort( + root, + model, + crate::ProviderKind::OpenRouter, + credentials, + None, + ) + .unwrap(); + Runtime::with_model_override(runtime, model_override).unwrap() + } + + async fn start(runtime: &Arc, mut claim: SessionClaim) -> (String, AcpDriver) { + let id = claim.request.id.clone(); + let driver = runtime + .start_acp_driver_with_mcp( + AcpDriverContext { + activity: crate::protocols::acp::activity::SessionActivity::new(|_| Ok(())), + ask_user: None, + cwd: runtime.root().to_path_buf(), + additional_directories: vec![], + integration: Arc::new(Observer), + cancellation: CancellationController::new().handle(), + response_attempt_replacement: false, + }, + &mut claim, + None, + runtime.mcp.clone(), + ) + .await + .unwrap(); + claim.commit().unwrap(); + (id, driver) + } + + async fn load(runtime: &Arc, id: &str) -> AcpDriver { + start(runtime, runtime.claim_session_load(id).unwrap()) + .await + .1 + } + + fn model(driver: &AcpDriver) -> String { + driver.adapter.selection().unwrap().model + } + + #[tokio::test] + async fn selected_model_survives_a_changed_default() { + let root = tempfile::tempdir().unwrap(); + let original = runtime(root.path(), "first/model", false); + let (id, driver) = start(&original, original.claim_session().unwrap()).await; + drop(driver); + drop(original); + + let restarted = runtime(root.path(), "changed/default", false); + let driver = load(&restarted, &id).await; + assert_eq!(model(&driver), "first/model"); + assert_eq!(driver.model_fallback, None); + driver + .adapter + .select(crate::provider::ModelSelection::new( + crate::ProviderKind::OpenRouter, + "second/model", + )) + .unwrap(); + drop(driver); + drop(restarted); + + let restarted = runtime(root.path(), "changed/default", false); + assert_eq!(model(&load(&restarted, &id).await), "second/model"); + } + + #[tokio::test] + async fn explicit_model_overrides_and_replaces_the_recorded_model() { + let root = tempfile::tempdir().unwrap(); + let original = runtime(root.path(), "first/model", false); + let (id, driver) = start(&original, original.claim_session().unwrap()).await; + drop(driver); + drop(original); + + let explicit = runtime(root.path(), "explicit/model", true); + let driver = load(&explicit, &id).await; + assert_eq!(model(&driver), "explicit/model"); + assert_eq!(driver.model_fallback, None); + assert!(!driver.unattributed_usage_applies); + drop(driver); + drop(explicit); + + let restarted = runtime(root.path(), "changed/default", false); + let driver = load(&restarted, &id).await; + assert_eq!(model(&driver), "explicit/model"); + assert!(!driver.unattributed_usage_applies); + } + + #[tokio::test] + async fn prompt_resume_with_an_explicit_model_replaces_the_recorded_model() { + let root = tempfile::tempdir().unwrap(); + let id = crate::session::new_id(); + let opened = crate::session::open( + root.path(), + &id, + false, + false, + vec![agentkit_core::Item::text( + agentkit_core::ItemKind::System, + "system", + )], + ) + .unwrap(); + opened.observer.set_model("openrouter:first/model").unwrap(); + drop(opened); + let runtime = Runtime::with_session_provider_credentials_effort_and_openrouter_key( + root.path(), + "explicit/model", + crate::ProviderKind::OpenRouter, + SessionRequest { + id: id.clone(), + resume: true, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + let runtime = Runtime::with_model_override(runtime, true).unwrap(); + // The empty key stops the run after session setup, before any model request. + assert!(runtime.run_persistent("hello".into()).await.is_err()); + drop(runtime); + let reopened = crate::session::open(root.path(), &id, true, true, vec![]).unwrap(); + assert_eq!(reopened.model.as_deref(), Some("openrouter:explicit/model")); + } + + #[tokio::test] + async fn transcript_without_a_model_record_uses_the_default_without_a_notice() { + let root = tempfile::tempdir().unwrap(); + let id = crate::session::new_id(); + drop( + crate::session::open( + root.path(), + &id, + false, + false, + vec![agentkit_core::Item::text( + agentkit_core::ItemKind::System, + "system", + )], + ) + .unwrap(), + ); + let driver = load(&runtime(root.path(), "the/default", false), &id).await; + assert_eq!(model(&driver), "the/default"); + assert_eq!(driver.model_fallback, None); + assert!(driver.unattributed_usage_applies); + drop(driver); + let overridden = load(&runtime(root.path(), "explicit/model", true), &id).await; + assert!(!overridden.unattributed_usage_applies); + } + + #[tokio::test] + async fn unusable_recorded_model_falls_back_with_a_notice_and_stays_recorded() { + for recorded in ["anthropic:no-credentials", "retired-provider:model"] { + let root = tempfile::tempdir().unwrap(); + let id = crate::session::new_id(); + let opened = crate::session::open( + root.path(), + &id, + false, + false, + vec![agentkit_core::Item::text( + agentkit_core::ItemKind::System, + "system", + )], + ) + .unwrap(); + opened.observer.set_model(recorded).unwrap(); + drop(opened); + + let driver = load(&runtime(root.path(), "the/default", false), &id).await; + assert_eq!(model(&driver), "the/default"); + let notice = driver.model_fallback.clone().unwrap(); + assert!(notice.contains(recorded), "{notice}"); + assert!(!driver.unattributed_usage_applies); + drop(driver); + + let reopened = crate::session::open(root.path(), &id, true, false, vec![]).unwrap(); + assert_eq!(reopened.model.as_deref(), Some(recorded)); + } + } +} + mod reasoning_effort_persistence { use super::*; use crate::ReasoningEffort; diff --git a/src/session.rs b/src/session.rs index bd880cc4..5c92043d 100644 --- a/src/session.rs +++ b/src/session.rs @@ -30,6 +30,7 @@ pub(crate) use children::{ChildLifecycle, DurableChild}; const CHILD_SCHEMA_VERSION: u32 = 5; // Older readers must reject selection records rather than silently lose settings. const REASONING_SCHEMA_VERSION: u32 = 6; +const MODEL_SCHEMA_VERSION: u32 = 7; pub const SCHEMA_VERSION: u32 = 3; const REDIRECT_SCHEMA_VERSION: u32 = 4; const PREVIOUS_SCHEMA_VERSION: u32 = 2; @@ -59,6 +60,8 @@ struct Record { snapshot: Option, #[serde(default, skip_serializing_if = "Option::is_none")] reasoning_effort: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + model: Option, } #[derive(Default, Serialize, Deserialize)] @@ -159,6 +162,7 @@ pub struct OpenSession { pub observer: SessionObserver, pub(crate) children: Vec, pub(crate) reasoning_effort: Option>, + pub(crate) model: Option, } /// Read-only metadata for one durable session in a workspace. @@ -247,6 +251,7 @@ struct Writer { // that owner publishes success, or forever if submission is abandoned. publication: Option>, reasoning_effort: Option>, + model: Option, } struct SessionLock { @@ -581,6 +586,7 @@ pub(crate) fn clone_completed_in( let authority = select_authority(directory, &canonical_workspace(root), source)? .ok_or_else(|| format!("session {source:?} does not exist"))?; let reasoning_effort = authority.reasoning_effort; + let model = authority.model; let mut transcript = authority.items; crate::transcript::repair_unanswered_tool_calls(&mut transcript); crate::transcript::rebind_forked_transcript(&mut transcript, destination); @@ -599,6 +605,9 @@ pub(crate) fn clone_completed_in( if let Some(effort) = reasoning_effort { opened.observer.set_reasoning_effort(effort)?; } + if let Some(model) = model { + opened.observer.set_model(&model)?; + } opened.observer.commit_creation()?; drop(opened); Ok(()) @@ -828,6 +837,7 @@ fn open_with_initial_timestamps_in( .map(|a| LiveSessionInfo::from_items(&a.historical_items, &a.items)) .unwrap_or_default(); let reasoning_effort = authority.as_ref().and_then(|a| a.reasoning_effort); + let model = authority.as_ref().and_then(|a| a.model.clone()); let children = authority .as_ref() .map(|a| a.children.clone()) @@ -857,6 +867,7 @@ fn open_with_initial_timestamps_in( created: (!resume).then(|| CreatedTranscript::new(None, filesystem, None)), publication: None, reasoning_effort, + model, }; return finish_open( writer, @@ -902,6 +913,7 @@ fn open_with_initial_timestamps_in( .map(|a| LiveSessionInfo::from_items(&a.historical_items, &a.items)) .unwrap_or_default(); let reasoning_effort = authority.as_ref().and_then(|a| a.reasoning_effort); + let model = authority.as_ref().and_then(|a| a.model.clone()); let children = authority .as_ref() .map(|a| a.children.clone()) @@ -920,7 +932,7 @@ fn open_with_initial_timestamps_in( &workspace_root, &authority.items, &authority.children, - (title_seed.as_deref(), reasoning_effort), + (title_seed.as_deref(), reasoning_effort, model.as_deref()), )?; for legacy in authority.legacy_histories { let source_fs = migration_filesystem(&legacy, &lock, &migration_locks)?; @@ -1006,6 +1018,7 @@ fn open_with_initial_timestamps_in( created, publication: None, reasoning_effort, + model, }; finish_open( writer, @@ -1056,6 +1069,7 @@ fn finish_open( } Ok(OpenSession { reasoning_effort: writer.reasoning_effort, + model: writer.model.clone(), children, transcript, observer: SessionObserver(Arc::new(Mutex::new(writer))), @@ -1154,6 +1168,15 @@ impl SessionObserver { .set_reasoning_effort(effort) } + /// Records the session's model id without changing transcript items. + /// Follows the same best-effort persistence semantics as reasoning effort. + pub(crate) fn set_model(&self, model: &str) -> Result<(), String> { + self.0 + .lock() + .map_err(|_| "session transcript writer poisoned".to_string())? + .set_model(model) + } + /// Appends parent-owned child state under the transcript lease and generation lock. /// Unlike ordinary transcript output, child lifecycle writes require disk. /// No registry callback runs while this writer guard is held; poison is isolation. @@ -1293,12 +1316,40 @@ impl Writer { .map_or("default", crate::ReasoningEffort::as_str) .into(), ), + model: None, }; self.write_record(record, generation)?; self.reasoning_effort = Some(effort); Ok(()) } + fn set_model(&mut self, model: &str) -> Result<(), String> { + if self.model.as_deref() == Some(model) { + return Ok(()); + } + self.ensure_lock()?; + let generation = self + .generation + .checked_add(1) + .ok_or_else(|| "session generation overflowed".to_string())?; + let record = Record { + schema_version: MODEL_SCHEMA_VERSION, + session_id: self.session_id.clone(), + generation, + workspace_root: Some(self.workspace_root.clone()), + item: None, + replacement: None, + redirect: None, + child: None, + snapshot: None, + reasoning_effort: None, + model: Some(model.into()), + }; + self.write_record(record, generation)?; + self.model = Some(model.into()); + Ok(()) + } + fn persist_child(&mut self, child: &DurableChild) -> Result<(), String> { child.validate()?; self.ensure_lock()?; @@ -1319,6 +1370,7 @@ impl Writer { redirect: None, child: Some(child.clone()), reasoning_effort: None, + model: None, snapshot: None, }; self.write_record(record, generation)?; @@ -1371,6 +1423,7 @@ impl Writer { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; self.write_record(record, generation) @@ -1392,6 +1445,7 @@ impl Writer { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; self.write_record(record, generation) @@ -1626,10 +1680,13 @@ enum HistoryEntry { Settings, } +type HistorySelection = (Option>, Option); + struct TranscriptHistory { reasoning_effort: Option>, + model: Option, children: Vec, - ancestry: Vec<(HistoryEntry, Option>)>, + ancestry: Vec<(HistoryEntry, HistorySelection)>, items: Vec, generation: u64, states: Vec>, @@ -1704,6 +1761,7 @@ fn read_record_lines( let mut children = BTreeMap::new(); let mut ancestry = Vec::new(); let mut reasoning_effort = None; + let mut model = None; let mut expected = 1_u64; let mut states = Vec::new(); let mut redirect = None; @@ -1719,12 +1777,13 @@ fn read_record_lines( | REDIRECT_SCHEMA_VERSION | CHILD_SCHEMA_VERSION | REASONING_SCHEMA_VERSION + | MODEL_SCHEMA_VERSION ) { return Err(format!( "unsupported session schema version {} on line {} (Kit supports {})", record.schema_version, index + 1, - REASONING_SCHEMA_VERSION + MODEL_SCHEMA_VERSION )); } if record.session_id != session_id || record.generation != expected { @@ -1744,19 +1803,40 @@ fn read_record_lines( .as_deref() .map(crate::ReasoningEffort::from_id) .transpose()?; - if selection.is_some() { - if record.schema_version != REASONING_SCHEMA_VERSION + let selected_model = record.model; + if selected_model.as_deref().is_some_and(str::is_empty) { + return Err("invalid session model record".into()); + } + let selects = selection.is_some() || selected_model.is_some(); + if selects { + let version_allowed = match selected_model { + Some(_) => record.schema_version == MODEL_SCHEMA_VERSION, + None => matches!( + record.schema_version, + REASONING_SCHEMA_VERSION | MODEL_SCHEMA_VERSION + ), + }; + if !version_allowed || record.workspace_root.is_none() || record.item.is_some() || record.child.is_some() || record.redirect.is_some() { - return Err("invalid reasoning effort record".into()); + return Err("invalid session selection record".into()); + } + if selection.is_some() { + reasoning_effort = selection; } - reasoning_effort = selection; - } else if record.schema_version == REASONING_SCHEMA_VERSION { - return Err("reasoning effort record missing selection".into()); + if selected_model.is_some() { + model = selected_model; + } + } else if matches!( + record.schema_version, + REASONING_SCHEMA_VERSION | MODEL_SCHEMA_VERSION + ) { + return Err("session selection record missing selection".into()); } + let current_selection = || (reasoning_effort, model.clone()); match ( record.item, record.replacement, @@ -1766,19 +1846,22 @@ fn read_record_lines( ) { (Some(item), None, None, None, None) if record.schema_version <= SCHEMA_VERSION => { items.push(item.clone()); - ancestry.push((HistoryEntry::Item(item), reasoning_effort)); + ancestry.push((HistoryEntry::Item(item), current_selection())); } (None, Some(replacement), None, None, None) if matches!( record.schema_version, - PREVIOUS_SCHEMA_VERSION | SCHEMA_VERSION | REASONING_SCHEMA_VERSION + PREVIOUS_SCHEMA_VERSION + | SCHEMA_VERSION + | REASONING_SCHEMA_VERSION + | MODEL_SCHEMA_VERSION ) && !replacement.is_empty() => { if !items.is_empty() { states.push(items.clone()); } items = replacement.clone(); - ancestry.push((HistoryEntry::Replacement(replacement), reasoning_effort)); + ancestry.push((HistoryEntry::Replacement(replacement), current_selection())); } (None, None, Some(target), None, None) if record.schema_version == REDIRECT_SCHEMA_VERSION @@ -1792,12 +1875,12 @@ fn read_record_lines( { child.validate()?; children.insert(child.id.clone(), child.clone()); - ancestry.push((HistoryEntry::Child(Box::new(child)), reasoning_effort)); + ancestry.push((HistoryEntry::Child(Box::new(child)), current_selection())); } (None, None, None, None, Some(snapshot)) if matches!( record.schema_version, - CHILD_SCHEMA_VERSION | REASONING_SCHEMA_VERSION + CHILD_SCHEMA_VERSION | REASONING_SCHEMA_VERSION | MODEL_SCHEMA_VERSION ) && record.workspace_root.is_some() && !snapshot.replacement.is_empty() => { @@ -1819,10 +1902,10 @@ fn read_record_lines( } items = snapshot.replacement.clone(); children = next; - ancestry.push((HistoryEntry::Snapshot(snapshot), reasoning_effort)); + ancestry.push((HistoryEntry::Snapshot(snapshot), current_selection())); } - (None, None, None, None, None) if selection.is_some() => { - ancestry.push((HistoryEntry::Settings, reasoning_effort)); + (None, None, None, None, None) if selects => { + ancestry.push((HistoryEntry::Settings, current_selection())); } _ => { return Err(format!( @@ -1842,6 +1925,7 @@ fn read_record_lines( states.push(items.clone()); Ok(StoredTranscript::History(TranscriptHistory { reasoning_effort, + model, children: children.into_values().collect(), ancestry, items, @@ -2625,6 +2709,7 @@ fn list_ids_in_with(filesystem: &Fs, directory: &Path) -> Result, St struct Authority { reasoning_effort: Option>, + model: Option, children: Vec, items: Vec, historical_items: Vec>, @@ -2650,7 +2735,7 @@ fn history_descends_from(history: &TranscriptHistory, ancestor: &TranscriptHisto let mut children: BTreeMap<&str, &DurableChild> = BTreeMap::new(); // A match must hold at one committed record boundary. In particular, a // catalog title seed must never supply transcript ancestry for a snapshot. - for (entry, reasoning_effort) in &history.ancestry { + for (entry, (reasoning_effort, model)) in &history.ancestry { match entry { HistoryEntry::Settings => {} HistoryEntry::Item(item) => items.push(item), @@ -2678,6 +2763,7 @@ fn history_descends_from(history: &TranscriptHistory, ancestor: &TranscriptHisto || legacy_prefix && items.len() >= ancestor.items.len()) && items.iter().zip(&ancestor.items).all(|(a, b)| *a == b) && *reasoning_effort == ancestor.reasoning_effort + && *model == ancestor.model && children.len() == ancestor.children.len() && children .values() @@ -2861,7 +2947,8 @@ fn select_authority_with( (true, false) => *current = candidate, (false, true) => {} (true, true) - if candidate.history.reasoning_effort == current.history.reasoning_effort => + if candidate.history.reasoning_effort == current.history.reasoning_effort + && candidate.history.model == current.history.model => { if candidate.path == scoped { *current = candidate; @@ -2881,6 +2968,7 @@ fn select_authority_with( } Ok(authority.map(|candidate| Authority { reasoning_effort: candidate.history.reasoning_effort, + model: candidate.history.model, children: candidate.history.children, items: candidate.history.items, historical_items: candidate.history.states, @@ -3147,9 +3235,22 @@ fn establish_scoped_authority( root: &Path, items: &[Item], children: &[DurableChild], - metadata: (Option<&[Item]>, Option>), + metadata: ( + Option<&[Item]>, + Option>, + Option<&str>, + ), ) -> Result<(), String> { - let (title_seed, reasoning_effort) = metadata; + let (title_seed, reasoning_effort, model) = metadata; + let schema_version = |fallback| { + if model.is_some() { + MODEL_SCHEMA_VERSION + } else if reasoning_effort.is_some() { + REASONING_SCHEMA_VERSION + } else { + fallback + } + }; let exists = fs::best_effort_global() .try_exists(path) .map_err(|error| format!("could not inspect {}: {error}", path.display()))?; @@ -3157,6 +3258,7 @@ fn establish_scoped_authority( match read_records_direct(path, session_id)? { StoredTranscript::History(history) if history.reasoning_effort == reasoning_effort + && history.model.as_deref() == model && history.items == items && history.children == children && transcript_workspace(path, session_id)?.as_deref() == Some(root) => @@ -3188,11 +3290,7 @@ fn establish_scoped_authority( filesystem, path, &Record { - schema_version: if reasoning_effort.is_some() { - REASONING_SCHEMA_VERSION - } else { - CHILD_SCHEMA_VERSION - }, + schema_version: schema_version(CHILD_SCHEMA_VERSION), session_id: session_id.into(), generation, workspace_root: Some(root.to_path_buf()), @@ -3205,6 +3303,7 @@ fn establish_scoped_authority( .map_or("default", crate::ReasoningEffort::as_str) .into() }), + model: model.map(str::to_owned), snapshot: Some(ChildSnapshot { replacement: items.to_vec(), children: children.to_vec(), @@ -3220,11 +3319,7 @@ fn establish_scoped_authority( filesystem, path, &Record { - schema_version: if reasoning_effort.is_some() { - REASONING_SCHEMA_VERSION - } else { - SCHEMA_VERSION - }, + schema_version: schema_version(SCHEMA_VERSION), session_id: session_id.into(), generation, workspace_root: Some(root.to_path_buf()), @@ -3237,6 +3332,7 @@ fn establish_scoped_authority( .map_or("default", crate::ReasoningEffort::as_str) .into() }), + model: model.map(str::to_owned), snapshot: None, }, create, @@ -3250,11 +3346,7 @@ fn establish_scoped_authority( filesystem, path, &Record { - schema_version: if reasoning_effort.is_some() { - REASONING_SCHEMA_VERSION - } else { - SCHEMA_VERSION - }, + schema_version: schema_version(SCHEMA_VERSION), session_id: session_id.into(), generation, workspace_root: Some(root.to_path_buf()), @@ -3267,6 +3359,7 @@ fn establish_scoped_authority( .map_or("default", crate::ReasoningEffort::as_str) .into() }), + model: model.map(str::to_owned), snapshot: None, }, create, @@ -3307,6 +3400,7 @@ fn redirect_legacy_transcript( redirect: Some(target), child: None, reasoning_effort: None, + model: None, snapshot: None, }, false, @@ -3669,6 +3763,124 @@ mod tests { } } + #[test] + fn model_survives_resume_and_fork() { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "selected", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + assert_eq!(opened.model, None); + let original = opened.transcript.clone(); + opened.observer.set_model("anthropic:first").unwrap(); + opened.observer.set_model("anthropic:second").unwrap(); + drop(opened); + let resumed = open(root.path(), "selected", true, false, vec![]).unwrap(); + assert_eq!(resumed.model.as_deref(), Some("anthropic:second")); + assert_eq!(resumed.transcript, original); + drop(resumed); + clone_completed(root.path(), "selected", "fork").unwrap(); + let fork = open(root.path(), "fork", true, false, vec![]).unwrap(); + assert_eq!(fork.model.as_deref(), Some("anthropic:second")); + } + + #[test] + fn unchanged_model_writes_no_record() { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "selected", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + opened.observer.set_model("anthropic:same").unwrap(); + opened.observer.checkpoint().unwrap(); + let path = transcript_path(root.path(), "selected"); + let before = fs::read_to_string(&path).unwrap(); + opened.observer.set_model("anthropic:same").unwrap(); + opened.observer.checkpoint().unwrap(); + assert_eq!(fs::read_to_string(&path).unwrap(), before); + } + + #[test] + fn model_records_validate_schema_and_payload() { + let base = json!({"schema_version": SCHEMA_VERSION, "session_id": "abc", + "generation": 1, "replacement": [Item::text(ItemKind::System, "system")]}); + let parse = |record: serde_json::Value| { + let bytes = format!("{base}\n{record}\n"); + read_records_bytes(Path::new("abc.jsonl"), "abc", bytes.as_bytes()) + }; + for version in [ + SCHEMA_VERSION, + CHILD_SCHEMA_VERSION, + REASONING_SCHEMA_VERSION, + ] { + assert!( + parse( + json!({"schema_version": version, "session_id": "abc", "generation": 2, + "workspace_root": "/workspace", "model": "anthropic:x"}) + ) + .is_err() + ); + } + for invalid in [ + json!({"schema_version": MODEL_SCHEMA_VERSION, "session_id": "abc", "generation": 2, + "workspace_root": "/workspace", "model": ""}), + json!({"schema_version": MODEL_SCHEMA_VERSION, "session_id": "abc", "generation": 2, + "workspace_root": "/workspace"}), + json!({"schema_version": MODEL_SCHEMA_VERSION, "session_id": "abc", "generation": 2, + "model": "anthropic:x"}), + ] { + assert!(parse(invalid).is_err()); + } + // An id from a removed provider must not make the transcript unreadable. + let Ok(StoredTranscript::History(history)) = parse( + json!({"schema_version": MODEL_SCHEMA_VERSION, "session_id": "abc", + "generation": 2, "workspace_root": "/workspace", "model": "retired:x"}), + ) else { + panic!("model record should parse"); + }; + assert_eq!(history.model.as_deref(), Some("retired:x")); + } + + #[test] + fn model_migration_preserves_selection_and_ancestry() { + let root = tempfile::tempdir().unwrap(); + let path = root.path().join("abc.jsonl"); + let items = vec![Item::text(ItemKind::System, "system")]; + let lock = SessionLock::acquire(path.with_extension("lock"), true).unwrap(); + let filesystem = lock.filesystem().unwrap(); + let establish = |model| { + establish_scoped_authority( + &filesystem, + &path, + "abc", + Path::new("/workspace"), + &items, + &[], + (None, None, model), + ) + .unwrap(); + let StoredTranscript::History(history) = read_records_direct(&path, "abc").unwrap() + else { + panic!() + }; + history + }; + let first = establish(Some("anthropic:first")); + assert_eq!(first.model.as_deref(), Some("anthropic:first")); + let second = establish(Some("anthropic:second")); + assert_eq!(second.model.as_deref(), Some("anthropic:second")); + assert!(history_descends_from(&second, &first)); + assert!(!history_descends_from(&first, &second)); + } + #[test] fn reasoning_effort_migration_preserves_joint_ancestry() { let root = tempfile::tempdir().unwrap(); @@ -3683,7 +3895,7 @@ mod tests { Path::new("/workspace"), &items, &[], - (None, Some(Some(crate::ReasoningEffort::High))), + (None, Some(Some(crate::ReasoningEffort::High)), None), ) .unwrap(); let StoredTranscript::History(high) = read_records_direct(&path, "abc").unwrap() else { @@ -3700,7 +3912,7 @@ mod tests { Path::new("/workspace"), &items, &[], - (None, Some(None)), + (None, Some(None), None), ) .unwrap(); let StoredTranscript::History(reset) = read_records_direct(&path, "abc").unwrap() else { @@ -3716,7 +3928,7 @@ mod tests { Path::new("/workspace"), &items, &[], - (None, Some(Some(crate::ReasoningEffort::High))), + (None, Some(Some(crate::ReasoningEffort::High)), None), ) .unwrap(); let StoredTranscript::History(high_again) = read_records_direct(&path, "abc").unwrap() @@ -3997,6 +4209,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }) .unwrap() @@ -4145,7 +4358,7 @@ mod tests { Path::new("/workspace"), &items, &authority.children, - (None, None), + (None, None, None), ) .unwrap(); let complete = fs::read_to_string(&path).unwrap(); @@ -4179,7 +4392,7 @@ mod tests { Path::new("/workspace"), &items, &authority.children, - (None, None), + (None, None, None), ) .unwrap(); let restored = parsed_child_history(&fs::read(&path).unwrap()); @@ -4433,6 +4646,7 @@ mod tests { redirect: None, child: Some(child.clone()), reasoning_effort: None, + model: None, snapshot: None, }, false, @@ -5085,6 +5299,7 @@ mod tests { redirect: Some(target), child: None, reasoning_effort: None, + model: None, snapshot: None, }; let mut contents = fs::read_to_string(&path).unwrap(); @@ -5147,6 +5362,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; fs::write( @@ -5280,6 +5496,7 @@ mod tests { redirect: Some(root.path().join("scoped/abc.jsonl")), child: None, reasoning_effort: None, + model: None, snapshot: None, }; @@ -5383,7 +5600,7 @@ mod tests { &workspace, &expected, &[], - (Some(&expected[..2]), None), + (Some(&expected[..2]), None, None), ) .unwrap(); drop(scoped_lock); @@ -5450,6 +5667,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }, false, @@ -5483,6 +5701,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let encoded = serde_json::to_vec(&record).unwrap(); @@ -5524,6 +5743,7 @@ mod tests { redirect: Some(normalized_absolute(&scoped).unwrap()), child: None, reasoning_effort: None, + model: None, snapshot: None, }; let encoded = serde_json::to_vec(&record).unwrap(); @@ -5575,6 +5795,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }, false, @@ -5693,6 +5914,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; fs::write( @@ -6495,6 +6717,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let mut encoded = serde_json::to_vec(&record).unwrap(); @@ -6921,6 +7144,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let original_user = Record { @@ -6933,6 +7157,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let current = vec![ @@ -6949,6 +7174,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let legacy = super::transcript_path(storage.path(), id); @@ -7035,6 +7261,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; fs::write( @@ -7091,6 +7318,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }, Record { @@ -7103,6 +7331,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }, ]; @@ -7154,6 +7383,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; fs::write( @@ -7296,6 +7526,7 @@ mod tests { redirect: None, child: None, reasoning_effort: None, + model: None, snapshot: None, }; let unscoped = super::transcript_path(storage.path(), "legacy-id"); diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 2e5d747d..35a3060a 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -964,6 +964,7 @@ fn agent_command_for_launch( session_id: &str, resume: bool, force: bool, + model_override: bool, ) -> std::io::Result { let mut command = crate::acp_child::serve_command( root, @@ -985,6 +986,10 @@ fn agent_command_for_launch( if force { command.arg("--force"); } + // The TUI always forwards its resolved model, so the server cannot infer intent. + if !model_override { + command.arg("--prefer-session-model"); + } Ok(command) } @@ -1375,6 +1380,7 @@ pub async fn run_with_reasoning_effort( resume, force, false, + false, &mut Stop::new()?, ) .await @@ -1395,6 +1401,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( resume: Option<&str>, force: bool, voice_enabled: bool, + model_override: bool, stop: &mut Stop, ) -> Result<(), Box> { run_client( @@ -1410,6 +1417,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( resume, force, voice_enabled, + model_override, stop, None, ) @@ -1475,6 +1483,7 @@ pub async fn run_remote( None, false, false, + false, stop, Some(remote), ) @@ -1496,6 +1505,7 @@ async fn run_client( resume: Option<&str>, force: bool, voice_enabled: bool, + model_override: bool, stop: &mut Stop, remote: Option<&crate::gateway::Remote>, ) -> Result<(), Box> { @@ -1555,6 +1565,7 @@ async fn run_client( launch_session_id, resume_session_id.is_some(), force, + model_override, )? }; let child_mcp_config = mcp_config.clone(); @@ -4847,9 +4858,15 @@ mod tests { "B", true, false, + false, ) .unwrap(); let resumed_args = command_args(&resumed); + assert!( + resumed_args + .iter() + .any(|arg| arg == "--prefer-session-model") + ); assert!( resumed_args .windows(2) @@ -4884,9 +4901,11 @@ mod tests { "C", false, false, + true, ) .unwrap(); let new_args = command_args(&new_session); + assert!(!new_args.iter().any(|arg| arg == "--prefer-session-model")); assert!( new_args .windows(2) diff --git a/tests/session_model.rs b/tests/session_model.rs new file mode 100644 index 00000000..81a8dcf8 --- /dev/null +++ b/tests/session_model.rs @@ -0,0 +1,177 @@ +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::disallowed_methods, + clippy::disallowed_macros +)] + +use std::{ + path::Path, + process::Stdio, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +#[derive(Default)] +struct Provider { + listed: Mutex>, + completions: Mutex>, +} + +async fn serve(listener: tokio::net::TcpListener, provider: Arc) { + loop { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0; 4096]; + let body = loop { + let n = stream.read(&mut buffer).await.unwrap(); + assert!(n > 0); + request.extend_from_slice(&buffer[..n]); + let Some(end) = request.windows(4).position(|s| s == b"\r\n\r\n") else { + continue; + }; + let headers = String::from_utf8_lossy(&request[..end]).to_string(); + let length: usize = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse().unwrap()) + }) + .unwrap_or(0); + if request.len() >= end + 4 + length { + break request[end + 4..end + 4 + length].to_vec(); + } + }; + let response = if request.starts_with(b"GET") { + let models: Vec<_> = provider + .listed + .lock() + .unwrap() + .iter() + .map(|id| serde_json::json!({"id": id})) + .collect(); + let body = serde_json::json!({"data": models}).to_string(); + format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + } else { + let request: serde_json::Value = serde_json::from_slice(&body).unwrap(); + provider + .completions + .lock() + .unwrap() + .push(request["model"].as_str().unwrap().to_owned()); + let body = "data: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2,\"total_tokens\":12}}\n\ndata: [DONE]\n\n"; + format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + }; + stream.write_all(response.as_bytes()).await.unwrap(); + } +} + +async fn prompt(home: &Path, url: &str, args: &[&str]) { + let mut command = tokio::process::Command::new(env!("CARGO_BIN_EXE_kit")); + command + .env_clear() + .env("HOME", home) + .env("PATH", "/usr/bin:/bin") + .env("OPENROUTER_API_KEY", "local-test-key") + .env("OPENROUTER_BASE_URL", url) + .args(["prompt", "--credential-store", "memory", "--root"]) + .arg(home) + .args(args) + .arg("hello") + .stdin(Stdio::null()) + .kill_on_drop(true); + let output = tokio::time::timeout(Duration::from_secs(30), command.output()) + .await + .expect("local prompt must settle") + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); +} + +fn transcript(home: &Path) -> std::path::PathBuf { + let paths: Vec<_> = std::fs::read_dir(home.join(".kit/sessions")) + .unwrap() + .flat_map(|workspace| std::fs::read_dir(workspace.unwrap().path()).unwrap()) + .map(|entry| entry.unwrap().path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "jsonl")) + .collect(); + assert_eq!(paths.len(), 1); + paths.into_iter().next().unwrap() +} + +fn latest_usage_model(path: &Path) -> serde_json::Value { + std::fs::read_to_string(path) + .unwrap() + .lines() + .rev() + .map(|line| serde_json::from_str::(line).unwrap()) + .find(|record| !record["item"]["usage"].is_null()) + .expect("transcript should record usage")["item"]["usage"]["metadata"]["kit.model"] + .clone() +} + +#[tokio::test] +async fn resumed_prompt_restores_overrides_and_falls_back_from_the_recorded_model() { + let home = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(home.path().join(".kit")).unwrap(); + std::fs::write( + home.path().join(".kit/config.toml"), + "provider = \"openrouter\"\nmodel = \"test/default\"\n", + ) + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!( + "http://{}/api/v1/chat/completions", + listener.local_addr().unwrap() + ); + let provider = Arc::new(Provider::default()); + *provider.listed.lock().unwrap() = vec!["test/default", "test/first", "test/explicit"]; + let server = tokio::spawn(serve(listener, Arc::clone(&provider))); + + prompt(home.path(), &url, &["--model", "test/first"]).await; + let path = transcript(home.path()); + let id = path.file_stem().unwrap().to_str().unwrap().to_owned(); + assert_eq!(latest_usage_model(&path), "openrouter:test/first"); + + prompt(home.path(), &url, &["--resume", &id]).await; + prompt( + home.path(), + &url, + &["--resume", &id, "--model", "test/explicit"], + ) + .await; + prompt(home.path(), &url, &["--resume", &id]).await; + assert_eq!(latest_usage_model(&path), "openrouter:test/explicit"); + + *provider.listed.lock().unwrap() = vec!["test/default"]; + prompt(home.path(), &url, &["--resume", &id]).await; + assert_eq!(latest_usage_model(&path), "openrouter:test/default"); + + *provider.listed.lock().unwrap() = vec!["test/default", "test/explicit"]; + prompt(home.path(), &url, &["--resume", &id]).await; + assert_eq!(latest_usage_model(&path), "openrouter:test/explicit"); + + server.abort(); + assert_eq!( + *provider.completions.lock().unwrap(), + [ + "test/first", + "test/first", + "test/explicit", + "test/explicit", + "test/default", + "test/explicit", + ] + ); +}