diff --git a/docs/user/tui-and-sessions.md b/docs/user/tui-and-sessions.md index 1dee5df..9f8f615 100644 --- a/docs/user/tui-and-sessions.md +++ b/docs/user/tui-and-sessions.md @@ -80,6 +80,7 @@ Press `F3` to open the read-only transcript navigator, even while a response is | `Ctrl+Up` / `Ctrl+Down` | Select the previous / next User prompt, switching to the User filter while retaining the query | | `Enter` | Reveal and highlight the selected rendered block in the transcript | | `Esc` / `F3` | Close without changing transcript scroll position, editor text, or attachments | +| Type `/branch`, then `Enter` | While idle, open the backend-approved text prompt checkout chooser | Navigator queries are limited to 4,096 UTF-8 bytes. Oversized pastes keep a bounded whole-grapheme prefix; pasted newlines and tabs are ignored rather than activating navigator controls. @@ -91,6 +92,24 @@ Revealing a Thought block temporarily shows it even when reasoning is hidden. Br The navigator searches only the currently displayed or replayed history, not a compacted archive or undelivered messages in the pending queue. Tool searches include display text such as the title, script, and output; media labels are searchable, but binary media payloads are not. Block identities are local and ephemeral, not durable addresses for forking. This is navigation only: it does not fork a session, write history, cancel a turn, or send a prompt. +### Edit a previous prompt in a new session + +While idle, type the exact local command `/branch`, or open `/transcript` (`F3`), type `/branch` in its search field, and press Enter. The separate checkout chooser lists only text prompts approved by the backend. Archived prompts are labeled `[archived]`; this list is authoritative, not inferred from the visible transcript. Use Up/Down to select a prompt and Enter to prepare an edit. Unsupported agents or ineligible prompts produce an error without changing the source session. + +> Only conversation context changes. Filesystem changes, running processes, and external effects are not rolled back. + +The TUI displays the backend's conversation prefix and puts the original prompt text in a **provisional prompt checkout** editor. Your source transcript, unsent editor draft, attachments, and model/configuration state are parked, not discarded. Edit the text, use Shift+Enter for a newline, then Enter to create and activate a new persisted session. No branch is created by merely browsing or preparing an edit. The edited prompt is persisted by branch submission itself; the TUI does not send it a second time as an ordinary prompt. + +Press Esc before submission to abandon the checkout and restore the parked source view, draft, attachments, and configuration. Esc also cancels a pending list or prepare request; late responses cannot replace a newer view. Once submission is in progress, wait for its result: Esc cannot undo a committed branch. A failed submit keeps the provisional draft available. Retry with the same text to recover a child if the response was lost. Once a submission reaches the backend, its checkout token is bound to that exact text; to submit a different edit, abandon and prepare a new checkout. + +Checkout edits are text-only. Adding image or audio attachments is rejected; attachments already in the parked source draft remain intact. While the checkout chooser or provisional editor is active, model/configuration changes, session switching, ordinary sends, and steering are disabled. Slash-command text in the provisional editor is edited prompt text, not a local command. + +The child retains the conversation strictly before the selected prompt, including its original bootstrap context. Archived prompts use validated pre-compaction history, never a later summary as a substitute for missing context. Prompts with unsupported content, unresolved tool calls in their prefix, or unreconstructable legacy context are not eligible. + +An unsubmitted checkout becomes stale when its source conversation or configuration changes, including compaction, or when the backend restarts. Abandon it and list prompts again. Committed submissions survive restart: retrying the same checkout and text finds the same child without generating a second response. Errors after a durable commit identify the child so it remains discoverable even if activation or response delivery failed. + +After successful submission, the source remains loaded and unchanged. Use `/sessions` to return to it. If the child is cancelled before execution starts, its committed history remains available but its connection can close to prevent the cancelled prompt from running later. Select the source, then the child in `/sessions` to reload it without rerunning that prompt. Checkout does not restore files, stop processes, reverse tool calls, or undo any other external effect. + ### Edit or remove a pending message Press `F2` to focus the queue, then `Up` or `Down` to select a message. The queue opens only when at least one message is pending; otherwise Kit shows `no pending messages` and leaves you in the composer. The selected row stays visible even when the queue is long. Press `Enter` to edit its text in the composer, or `Backspace` / `Delete` to request removal (the normal Mac Delete key works). `Esc` or `F2` returns from the selector without changing your draft. When the last message is delivered or removed, or the turn finishes, the selector closes automatically and keyboard focus returns to the composer. An in-progress text edit remains available to copy or cancel rather than being discarded. @@ -139,7 +158,7 @@ When the agent roster is visible, terminals at least 108 columns wide show the t ## Manage sessions and compact from the TUI -The TUI handles `/new`, `/resume`, `/sessions`, `/close`, `/model`, `/effort`, `/agents`, and `/transcript` as exact local slash-command tokens. It also discovers agent commands through ACP and highlights them without interpreting them locally: +The TUI handles `/new`, `/resume`, `/sessions`, `/close`, `/model`, `/effort`, `/agents`, `/transcript`, and `/branch` as exact local slash-command tokens. It also discovers agent commands through ACP and highlights them without interpreting them locally: ```text /new diff --git a/src/compaction.rs b/src/compaction.rs index 983a7cc..00d2c30 100644 --- a/src/compaction.rs +++ b/src/compaction.rs @@ -571,6 +571,8 @@ impl CompactionStrategy for SummarizeForContinuation { .rev() .find(|item| is_compaction_summary(item)) .cloned(); + // Keep bootstrap Items intact: their metadata includes versioned branch + // lineage and captured model selection needed to resume after compaction. let bootstrap = request .transcript .iter() @@ -877,6 +879,59 @@ mod tests { assert_eq!(text.text, "current request"); } + #[tokio::test] + async fn strategy_preserves_prompt_checkout_lineage_and_selection() { + use crate::session::branch::{ + Boundary, BranchMetadata, CapturedSelection, SubmittedRequest, prepare, + }; + + let prefix = vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::Context, "AGENTS"), + Item::text(ItemKind::User, "old request"), + Item::text(ItemKind::Assistant, "old answer"), + ]; + let selection = CapturedSelection::new( + &crate::provider::ModelSelection::new( + crate::ProviderKind::OpenRouter, + "openai/gpt-5.4", + ), + Some(crate::ReasoningEffort::High), + ); + let mut transcript = prepare( + prefix.clone(), + "parent".into(), + Boundary::new(0, &prefix).unwrap(), + "checkout".into(), + SubmittedRequest { + id: "request".into(), + selection: selection.clone(), + }, + Item::text(ItemKind::User, "edited request"), + ) + .unwrap(); + let bootstrap = transcript[0].clone(); + let lineage = BranchMetadata::read(&transcript).unwrap().unwrap(); + let backend = FixedBackend; + // Exercise both summarized and recent-only retention, then compact the + // result again. Lineage belongs to the unchanged bootstrap, not summary. + for recent_tokens in [1, usize::MAX, 1] { + let mut context = CompactionContext::new().with_backend(&backend); + transcript = SummarizeForContinuation { recent_tokens } + .apply( + CompactionRequest::new(transcript, CompactionReason::TranscriptTooLong), + &mut context, + ) + .await + .unwrap() + .transcript; + assert_eq!(transcript[0], bootstrap); + let restored = BranchMetadata::read(&transcript).unwrap().unwrap(); + assert_eq!(restored, lineage); + assert_eq!(restored.request.selection, selection); + } + } + #[tokio::test] async fn strategy_keeps_recent_tool_round_at_tail_during_mid_turn_compaction() { let transcript = vec![ diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index e6e4b47..09909b3 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -49,6 +49,7 @@ use tracing::Instrument as _; mod activity; pub(crate) mod model_switch; +pub(crate) mod prompt_branches; mod skill_catalog; pub mod v2; @@ -1474,7 +1475,7 @@ impl Server { runtime: Arc::clone(&self.runtime), integration: Arc::clone(&self.integration), binding, - driver: driver.driver, + driver: driver.driver.into_inner(), tasks: driver.tasks, background_jobs: background_jobs.clone(), structured_completion, diff --git a/src/protocols/acp/prompt_branches.rs b/src/protocols/acp/prompt_branches.rs new file mode 100644 index 0000000..f038804 --- /dev/null +++ b/src/protocols/acp/prompt_branches.rs @@ -0,0 +1,629 @@ +//! Kit-private text-prompt checkout protocol. Addresses are issued by the backend; +//! neither display positions nor provider item identifiers are branch authority. + +use agentkit_acp::v2::wire; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcRequest)] +#[request(method = "kit/transcript/list", response = ListPromptBranchesResponse)] +#[serde(deny_unknown_fields)] +pub(crate) struct ListPromptBranchesRequest { + pub session_id: wire::SessionId, +} + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcResponse)] +pub(crate) struct ListPromptBranchesResponse { + pub boundaries: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct PromptBoundary { + pub address: String, + pub text: String, + pub historical: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcRequest)] +#[request(method = "kit/transcript/prepare", response = PreparePromptBranchResponse)] +#[serde(deny_unknown_fields)] +pub(crate) struct PreparePromptBranchRequest { + pub session_id: wire::SessionId, + pub address: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcResponse)] +pub(crate) struct PreparePromptBranchResponse { + pub checkout_token: String, + pub original_text: String, + pub prefix: Vec, + pub config_options: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcRequest)] +#[request(method = "kit/transcript/submit", response = SubmitPromptBranchResponse)] +#[serde(deny_unknown_fields)] +pub(crate) struct SubmitPromptBranchRequest { + pub session_id: wire::SessionId, + pub checkout_token: String, + pub text: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcResponse)] +pub(crate) struct SubmitPromptBranchResponse { + pub session_id: wire::SessionId, + pub config_options: Vec, +} + +use std::{ + collections::HashMap, + path::Path, + sync::{Arc, Mutex}, +}; + +use crate::{ + provider::{ModelSelection, ReasoningEffort}, + runtime::AcpForkState, + session::branch, +}; +use agentkit_core::{Item, ItemKind, Part}; + +/// Actor-local authority. Dropping the actor invalidates every uncommitted +/// address and checkout, even if the same durable session is loaded again. +#[derive(Default)] +pub(crate) struct PromptCheckouts { + boundaries: HashMap, + checkouts: HashMap, +} + +#[derive(Clone)] +struct SourceRevision { + transcript: Arc>, + // Monotonic actor revision for transcript work and successful config edits. + // Exact transcript equality alone cannot detect mutation followed by compaction. + checkout_revision: u64, + selection: ModelSelection, + reasoning: Option, +} + +impl SourceRevision { + fn validate( + &self, + transcript: &[Item], + checkout_revision: u64, + selection: &ModelSelection, + reasoning: Option, + ) -> Result<(), String> { + if self.transcript.as_slice() != transcript + || self.checkout_revision != checkout_revision + || &self.selection != selection + || self.reasoning != reasoning + { + return Err( + "stale prompt checkout: transcript or configuration changed; list prompts again" + .into(), + ); + } + Ok(()) + } +} + +#[derive(Clone)] +struct ApprovedBoundary { + source: SourceRevision, + parent_session_id: String, + provenance: branch::Boundary, + state: Arc>, + text: String, +} + +#[derive(Clone)] +pub(crate) struct PreparedCheckout { + pub token: String, + pub original_text: String, + pub prefix: Vec, + pub selection: ModelSelection, + pub reasoning: Option, + boundary: ApprovedBoundary, + submitted: Arc>>, +} + +impl PreparedCheckout { + /// Construct only in memory. Runtime's guarded initialization owns the + /// completion record and disk barrier, before adapter startup or execution. + pub(crate) fn fork(&self, text: &str) -> Result { + if text.trim().is_empty() { + return Err("prompt checkout requires non-empty text".into()); + } + let request_id = submitted_request_id(&self.boundary.parent_session_id, text); + let mut submitted = self + .submitted + .lock() + .map_err(|_| "prompt checkout identity lock poisoned")?; + if submitted + .as_ref() + .is_some_and(|previous| previous != &request_id) + { + return Err( + "this checkout token is already bound to a different submitted request".into(), + ); + } + *submitted = Some(request_id.clone()); + let mut prefix = self.prefix.clone(); + crate::transcript::sanitize_forked_transcript(&mut prefix); + let transcript = branch::prepare( + prefix, + self.boundary.parent_session_id.clone(), + self.boundary.provenance.clone(), + self.token.clone(), + branch::SubmittedRequest { + id: request_id, + selection: branch::CapturedSelection::new(&self.selection, self.reasoning), + }, + Item::text(ItemKind::User, text), + )?; + Ok(AcpForkState { + transcript, + selection: self.selection.clone(), + reasoning_effort: self.reasoning, + parent_context: None, + }) + } +} + +impl PromptCheckouts { + #[allow(clippy::too_many_arguments)] + pub(crate) fn list( + &mut self, + root: &Path, + session_id: &str, + transcript: &[Item], + checkout_revision: u64, + selection: &ModelSelection, + reasoning: Option, + ) -> Result { + let states = branch::load_history(root, session_id)?; + self.list_states( + session_id, + transcript, + &states, + checkout_revision, + selection, + reasoning, + ) + } + + #[allow(clippy::too_many_arguments)] + fn list_states( + &mut self, + session_id: &str, + transcript: &[Item], + states: &[Vec], + checkout_revision: u64, + selection: &ModelSelection, + reasoning: Option, + ) -> Result { + // The historical loader must resolve to the actor's current canonical + // state. Never replace missing old context with today's compacted state. + if states.last().map(Vec::as_slice) != Some(transcript) { + return Err("session history does not match the loaded transcript; reload and list prompts again".into()); + } + let source = SourceRevision { + transcript: Arc::new(transcript.to_vec()), + checkout_revision, + selection: selection.clone(), + reasoning, + }; + self.boundaries.retain(|_, boundary| { + boundary + .source + .validate(transcript, checkout_revision, selection, reasoning) + .is_ok() + }); + self.checkouts.retain(|_, checkout| { + checkout + .boundary + .source + .validate(transcript, checkout_revision, selection, reasoning) + .is_ok() + }); + let mut result = Vec::new(); + let mut seen = std::collections::HashSet::new(); + // Current state first: unchanged prompts retained across compaction are + // not duplicated as archived points. Every other prefix stays historical. + for (state_index, state) in states.iter().enumerate().rev() { + branch::BranchMetadata::read(state)?; + if state + .first() + .is_none_or(|item| item.kind != ItemKind::System) + { + continue; // A legacy transcript without bootstrap is unreconstructable. + } + let state = Arc::new(state.clone()); + for (index, item) in state.iter().enumerate() { + if item.kind != ItemKind::User { + continue; + } + let Ok(text) = text_prompt(item) else { + continue; + }; + let prefix = &state[..index]; + if prefix.is_empty() || crate::transcript::has_unanswered_tool_calls(prefix) { + continue; + } + let identity = + serde_json::to_vec(&state[..=index]).map_err(|error| error.to_string())?; + if !seen.insert(blake3::hash(&identity)) { + continue; + } + let provenance = branch::Boundary::new(state_index, prefix)?; + let address = self + .boundaries + .iter() + .find_map(|(address, existing)| { + (existing.provenance == provenance && existing.text == text) + .then(|| address.clone()) + }) + .unwrap_or_else(crate::session::new_id); + let boundary = ApprovedBoundary { + source: source.clone(), + parent_session_id: session_id.into(), + provenance, + state: Arc::clone(&state), + text: text.clone(), + }; + self.boundaries.insert(address.clone(), boundary); + result.push(PromptBoundary { + address, + text, + historical: state_index + 1 != states.len(), + }); + } + } + Ok(ListPromptBranchesResponse { boundaries: result }) + } + + pub(crate) fn prepare( + &mut self, + address: &str, + transcript: &[Item], + checkout_revision: u64, + selection: &ModelSelection, + reasoning: Option, + ) -> Result { + let boundary = self + .boundaries + .get(address) + .ok_or_else(|| "unknown or stale prompt address; list prompts again".to_string())?; + boundary + .source + .validate(transcript, checkout_revision, selection, reasoning)?; + let checkout = PreparedCheckout { + token: crate::session::new_id(), + original_text: boundary.text.clone(), + prefix: boundary.state[..boundary.provenance.prefix_len].to_vec(), + selection: selection.clone(), + reasoning, + boundary: boundary.clone(), + submitted: Arc::new(Mutex::new(None)), + }; + self.checkouts + .insert(checkout.token.clone(), checkout.clone()); + Ok(checkout) + } + + pub(crate) fn checkout( + &self, + token: &str, + transcript: &[Item], + checkout_revision: u64, + selection: &ModelSelection, + reasoning: Option, + ) -> Result { + let checkout = self + .checkouts + .get(token) + .ok_or_else(|| "unknown or stale checkout token; list prompts again".to_string())?; + checkout + .boundary + .source + .validate(transcript, checkout_revision, selection, reasoning)?; + Ok(checkout.clone()) + } +} + +fn text_prompt(item: &Item) -> Result { + if item.kind != ItemKind::User || item.parts.is_empty() { + return Err("checkout supports text-only user prompts".into()); + } + let mut text = String::new(); + for part in &item.parts { + let Part::Text(part) = part else { + return Err( + "checkout does not support media, resources, or other non-text prompt parts".into(), + ); + }; + text.push_str(&part.text); + } + if text.trim().is_empty() { + return Err("checkout requires non-empty prompt text".into()); + } + Ok(text) +} + +pub(crate) fn submitted_request_id(source_session_id: &str, text: &str) -> String { + let bytes = serde_json::to_vec(&(source_session_id, text)).expect("string tuples serialize"); + blake3::hash(&bytes).to_hex().to_string() +} + +/// Resolve durable completion before asking a source actor for admission. This +/// path also works after restart, when all in-memory addresses have expired. +pub(crate) fn lookup_committed( + root: &Path, + source_session_id: &str, + token: &str, + text: &str, +) -> Result, String> { + branch::find_committed( + root, + source_session_id, + token, + &submitted_request_id(source_session_id, text), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use agentkit_core::{DataRef, Modality, ToolCallPart}; + use serde_json::json; + + fn selection() -> ModelSelection { + ModelSelection::new(crate::ProviderKind::OpenRouter, "test/model") + } + + fn conversation() -> Vec { + vec![ + Item::text(ItemKind::System, "bootstrap"), + Item::text(ItemKind::Context, "workspace context"), + Item::text(ItemKind::User, "first"), + Item::text(ItemKind::Assistant, "original future"), + Item::text(ItemKind::User, "second"), + Item::text(ItemKind::Assistant, "second future"), + ] + } + + fn list( + checkouts: &mut PromptCheckouts, + current: &[Item], + states: &[Vec], + ) -> ListPromptBranchesResponse { + checkouts + .list_states("source", current, states, 0, &selection(), None) + .unwrap() + } + + #[test] + fn exclusive_prefix_preserves_first_bootstrap_and_original_future() { + let source = conversation(); + let before = source.clone(); + let mut checkouts = PromptCheckouts::default(); + let boundaries = list(&mut checkouts, &source, std::slice::from_ref(&source)); + assert_eq!(boundaries.boundaries.len(), 2); + for (index, prefix_len) in [2, 4].into_iter().enumerate() { + let prepared = checkouts + .prepare( + &boundaries.boundaries[index].address, + &source, + 0, + &selection(), + None, + ) + .unwrap(); + assert_eq!(prepared.prefix, source[..prefix_len]); + assert_eq!( + prepared.original_text, + if index == 0 { "first" } else { "second" } + ); + let fork = prepared.fork("edited").unwrap(); + assert_eq!(fork.transcript.len(), prefix_len + 1); + assert_eq!( + text_prompt(fork.transcript.last().unwrap()).unwrap(), + "edited" + ); + assert_eq!(&fork.transcript[1..prefix_len], &source[1..prefix_len]); + assert!( + branch::BranchMetadata::read(&fork.transcript) + .unwrap() + .is_some() + ); + } + assert_eq!(source, before); + } + + #[test] + fn historical_checkout_never_substitutes_future_summary() { + let historical = conversation(); + let mut summary = Item::text(ItemKind::Developer, "FUTURE SUMMARY of everything"); + summary + .metadata + .insert("kit.compaction.summary".into(), true.into()); + let current = vec![ + historical[0].clone(), + historical[1].clone(), + summary, + Item::text(ItemKind::User, "after compaction"), + ]; + let states = vec![historical.clone(), current.clone()]; + let mut checkouts = PromptCheckouts::default(); + let boundaries = list(&mut checkouts, ¤t, &states); + let old = boundaries + .boundaries + .iter() + .find(|b| b.text == "second") + .unwrap(); + assert!(old.historical); + let prepared = checkouts + .prepare(&old.address, ¤t, 0, &selection(), None) + .unwrap(); + assert_eq!(prepared.prefix, historical[..4]); + assert!( + !serde_json::to_string(&prepared.prefix) + .unwrap() + .contains("FUTURE SUMMARY") + ); + let now = boundaries + .boundaries + .iter() + .find(|b| b.text == "after compaction") + .unwrap(); + assert!(!now.historical); + } + + #[test] + fn read_only_lists_and_prepares_do_not_stale_existing_tokens() { + let source = conversation(); + let mut checkouts = PromptCheckouts::default(); + let first = list(&mut checkouts, &source, std::slice::from_ref(&source)); + let prepared = checkouts + .prepare(&first.boundaries[0].address, &source, 0, &selection(), None) + .unwrap(); + let again = list(&mut checkouts, &source, std::slice::from_ref(&source)); + assert_eq!(first.boundaries[0].address, again.boundaries[0].address); + checkouts + .prepare(&first.boundaries[1].address, &source, 0, &selection(), None) + .unwrap(); + assert!( + checkouts + .checkout(&prepared.token, &source, 0, &selection(), None) + .is_ok() + ); + assert_eq!(checkouts.boundaries.len(), 2); + // Restart creates a fresh actor-local authority regardless of Item.id. + assert!( + PromptCheckouts::default() + .checkout(&prepared.token, &source, 0, &selection(), None) + .is_err() + ); + assert!(source.iter().all(|item| item.id.is_none())); + } + + #[test] + fn transcript_checkout_revision_model_and_reasoning_changes_stale_checkout() { + let source = conversation(); + let mut checkouts = PromptCheckouts::default(); + let boundaries = list(&mut checkouts, &source, std::slice::from_ref(&source)); + let prepared = checkouts + .prepare( + &boundaries.boundaries[0].address, + &source, + 0, + &selection(), + None, + ) + .unwrap(); + let mut future = source.clone(); + future.push(Item::text(ItemKind::User, "new future")); + assert!( + checkouts + .checkout(&prepared.token, &future, 0, &selection(), None) + .is_err() + ); + assert!( + checkouts + .checkout(&prepared.token, &source, 1, &selection(), None) + .is_err() + ); + let different = ModelSelection::new(crate::ProviderKind::OpenRouter, "different/model"); + assert!( + checkouts + .checkout(&prepared.token, &source, 0, &different, None) + .is_err() + ); + assert!( + checkouts + .checkout( + &prepared.token, + &source, + 0, + &selection(), + Some(ReasoningEffort::High) + ) + .is_err() + ); + assert!( + checkouts + .prepare("0", &source, 0, &selection(), None) + .is_err() + ); + } + + #[test] + fn rejects_nontext_prompts_open_tool_prefix_and_missing_bootstrap() { + let media = Item::new( + ItemKind::User, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1]), + )], + ); + assert!( + text_prompt(&media) + .unwrap_err() + .contains("media, resources") + ); + let source = vec![ + Item::text(ItemKind::System, "bootstrap"), + media, + Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new("call", "tool", json!({})))], + ), + Item::text(ItemKind::User, "unresolved prefix"), + ]; + let mut checkouts = PromptCheckouts::default(); + assert!( + list(&mut checkouts, &source, std::slice::from_ref(&source)) + .boundaries + .is_empty() + ); + let legacy = vec![Item::text(ItemKind::User, "lost context")]; + assert!( + list(&mut checkouts, &legacy, std::slice::from_ref(&legacy)) + .boundaries + .is_empty() + ); + assert!( + checkouts + .list_states("source", &source, &[legacy], 0, &selection(), None) + .is_err() + ); + assert!(serde_json::from_value::(json!({ + "session_id":"source", "checkout_token":"token", "text":"hello", "attachments":["image"] + })).is_err()); + } + + #[test] + fn checkout_request_identity_survives_clones_and_rejects_different_edit() { + let source = conversation(); + let mut checkouts = PromptCheckouts::default(); + let boundaries = list(&mut checkouts, &source, std::slice::from_ref(&source)); + let prepared = checkouts + .prepare( + &boundaries.boundaries[0].address, + &source, + 0, + &selection(), + None, + ) + .unwrap(); + let retry = prepared.clone(); + assert!(prepared.fork("accepted").is_ok()); + assert!(retry.fork("accepted").is_ok()); + assert!(retry.fork("different").is_err()); + assert_ne!( + submitted_request_id("source", "x"), + submitted_request_id("different source", "x") + ); + } +} diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index d01d5ac..ad3673d 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, path::{Path, PathBuf}, sync::{ Arc, Mutex, Weak, @@ -20,21 +20,26 @@ use agentkit_core::{ ToolOutput, Usage, }; use agentkit_loop::{ - AgentEvent, LoopDriver, LoopError, LoopInterrupt, LoopObserver, LoopStep, ModelSession, - ObservedEvent, + AgentEvent, LoopError, LoopInterrupt, LoopObserver, LoopStep, ModelSession, ObservedEvent, }; use agentkit_task_manager::{TaskEvent, TaskManagerHandle}; use async_trait::async_trait; -use tokio::sync::{mpsc, oneshot, watch}; +use futures_util::FutureExt; +use tokio::sync::{Notify, mpsc, oneshot, watch}; use tracing::Instrument as _; use crate::{ provider::{ProviderKind, SelectableAdapter, authentication_method_id}, - runtime::{AcpDriverContext, BackgroundJobs, Runtime}, + runtime::{AcpDriverContext, BackgroundJobs, InputSettlingDriver as LoopDriver, Runtime}, }; use super::activity::{ExecutionOrigin, SessionActivity}; use super::model_switch; +use super::prompt_branches::{ + self, ListPromptBranchesRequest, ListPromptBranchesResponse, PreparePromptBranchRequest, + PreparePromptBranchResponse, PreparedCheckout, PromptCheckouts, SubmitPromptBranchRequest, + SubmitPromptBranchResponse, +}; use super::{ AuthenticationRequiredData, CancelBackgroundRequest, CancelBackgroundResponse, @@ -44,9 +49,33 @@ use super::{ const PAGE_SIZE: usize = 100; +static BRANCH_SUBMISSIONS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + static NEXT_ERROR_MESSAGE_ID: AtomicU64 = AtomicU64::new(1); static NEXT_THOUGHT_MESSAGE_ID: AtomicU64 = AtomicU64::new(1); +fn establish_diagnostic_route(session_id: &wire::SessionId) { + crate::events::emit(&crate::events::RuntimeEvent::SessionStarted { + session_id: session_id.to_string(), + }); +} + +fn validate_resume_location( + root: &std::path::Path, + cwd: &std::path::Path, + has_additional_directories: bool, +) -> Result<(), AcpRuntimeError> { + let cwd = crate::resilient_fs::canonicalize(cwd) + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; + if cwd != root || has_additional_directories { + return Err(AcpRuntimeError::Loop(format!( + "this Kit runtime is fixed to {} and does not accept additional directories", + root.display() + ))); + } + Ok(()) +} + fn available_commands_update(session_id: wire::SessionId) -> wire::UpdateSessionNotification { wire::UpdateSessionNotification::new( session_id, @@ -139,14 +168,97 @@ fn claim_prompt(busy: &AtomicBool) -> Result<(), AcpRuntimeError> { } #[derive(Clone)] -struct ConnectionSink(V2ConnectionTo); +struct ConnectionSink(V2ConnectionTo, Arc>); + +// The integration's pending queue is private. Track admission before awaiting +// its reservation, then track accepted IDs until delivery or successful revoke. +#[derive(Default)] +struct InjectionWork { + admitting: usize, + pending: HashSet, + branch_reserved: bool, + admission_released: Arc, +} + +struct InjectionAdmission(Arc>); + +impl Drop for InjectionAdmission { + fn drop(&mut self) { + self.0.lock().expect("injection work poisoned").admitting -= 1; + } +} + +struct TrackedInjection { + work: Arc>, + id: wire::MessageId, + retained: bool, +} + +impl Drop for TrackedInjection { + fn drop(&mut self) { + if !self.retained { + self.work + .lock() + .expect("injection work poisoned") + .pending + .remove(&self.id); + } + } +} + +struct BranchAdmission { + busy: Arc, + injections: Arc>, +} + +impl BranchAdmission { + fn claim( + busy: Arc, + injections: Arc>, + ) -> Result { + { + let mut work = injections.lock().expect("injection work poisoned"); + if work.branch_reserved || work.admitting != 0 || !work.pending.is_empty() { + return Err(AcpRuntimeError::Loop( + "source session has queued injection work".into(), + )); + } + claim_prompt(&busy)?; + work.branch_reserved = true; + } + Ok(Self { busy, injections }) + } +} + +impl Drop for BranchAdmission { + fn drop(&mut self) { + let mut work = self.injections.lock().expect("injection work poisoned"); + work.branch_reserved = false; + self.busy.store(false, Ordering::Release); + // Admission can be dropped before its command reaches the actor. Keep + // a permit so a selected autonomous event is retried even in that case. + work.admission_released.notify_one(); + } +} #[async_trait] impl AcpSessionUpdateSink for ConnectionSink { fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError> { + let delivered = match ¬ification.update { + wire::SessionUpdate::UserMessage(message) => Some(message.message_id.clone()), + _ => None, + }; self.0 .send_notification(notification) - .map_err(|error| AcpRuntimeError::Sdk(error.to_string())) + .map_err(|error| AcpRuntimeError::Sdk(error.to_string()))?; + if let Some(id) = delivered { + self.1 + .lock() + .expect("injection work poisoned") + .pending + .remove(&id); + } + Ok(()) } async fn update_acknowledged( @@ -207,10 +319,89 @@ impl ReplacementGeneration { } } +// A boundary owns this ledger while the SDK submits and acknowledges one +// injection at a time. Ordinary prompt echoes use `update`, not this ledger. +// Successful message batches retain their ordered input items even if a later +// notification fails. Display message IDs are not canonical Item identities. +#[derive(Clone, Default)] +struct SteeringAcknowledgements(Arc>>>); + +struct SteeringBoundary<'a>(&'a SteeringAcknowledgements); + +impl SteeringAcknowledgements { + fn begin(&self) -> Result, AcpRuntimeError> { + let mut state = self.0.lock().map_err(|_| AcpRuntimeError::ClientClosed)?; + if state.is_some() { + return Err(AcpRuntimeError::Loop( + "injection boundary already active".into(), + )); + } + *state = Some(Vec::new()); + Ok(SteeringBoundary(self)) + } + + fn prepare( + &self, + notification: &wire::UpdateSessionNotification, + ) -> Result>, AcpRuntimeError> { + let active = self + .0 + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)? + .is_some(); + if !active { + return Ok(None); + } + let wire::SessionUpdate::UserMessage(message) = ¬ification.update else { + return Ok(None); + }; + let agent_client_protocol::schema::MaybeUndefined::Value(content) = &message.content else { + return Err(AcpRuntimeError::Loop( + "injected user message has no content".into(), + )); + }; + agentkit_acp::v2::content_blocks_to_items(content).map(Some) + } + + fn acknowledged(&self, items: Option>) -> Result<(), AcpRuntimeError> { + if let Some(items) = items { + self.0 + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)? + .as_mut() + .ok_or_else(|| { + AcpRuntimeError::Loop("injection boundary ended before acknowledgement".into()) + })? + .extend(items); + } + Ok(()) + } +} + +impl SteeringBoundary<'_> { + fn finish(self) -> Result, AcpRuntimeError> { + let completed = { + let mut state = self.0.0.lock().map_err(|_| AcpRuntimeError::ClientClosed)?; + state.take() + }; + completed.ok_or_else(|| AcpRuntimeError::Loop("injection boundary is not active".into())) + } +} + +impl Drop for SteeringBoundary<'_> { + fn drop(&mut self) { + // No recovery of poisoned acknowledgement state. Abandon its owner; + // otherwise release the items outside the guard, including on unwind. + let abandoned = self.0.0.lock().ok().and_then(|mut state| state.take()); + drop(abandoned); + } +} + #[derive(Clone)] struct ResponseReplacementSink { inner: S, current: Arc>, + acknowledgements: SteeringAcknowledgements, } impl ResponseReplacementSink { @@ -218,6 +409,7 @@ impl ResponseReplacementSink { Self { inner, current: Arc::new(Mutex::new(CurrentReplacementMessages::default())), + acknowledgements: SteeringAcknowledgements::default(), } } @@ -356,8 +548,11 @@ impl AcpSessionUpdateSink for ResponseReplacementSink Result<(), AcpRuntimeError> { + let items = self.acknowledgements.prepare(¬ification)?; self.rewrite_and_track(&mut notification); - self.inner.update_acknowledged(notification).await + // Never hold the ledger guard across external delivery or its await. + self.inner.update_acknowledged(notification).await?; + self.acknowledgements.acknowledged(items) } async fn flush(&self) -> Result<(), AcpRuntimeError> { @@ -463,6 +658,22 @@ struct PromptCommand { } enum Command { + Snapshot { + reply: oneshot::Sender>, + }, + ListPromptBranches { + reply: oneshot::Sender>, + }, + PreparePromptBranch { + address: String, + admission: BranchAdmission, + reply: oneshot::Sender>, + }, + ReservePromptBranch { + checkout_token: String, + admission: BranchAdmission, + reply: oneshot::Sender), AcpRuntimeError>>, + }, Prompt(PromptCommand), SetConfig { request: wire::SetSessionConfigOptionRequest, @@ -477,6 +688,7 @@ enum Command { } struct SessionHandle { + injections: Arc>, token: u64, commands: mpsc::Sender, integration: AcpSessionHandle, @@ -486,6 +698,11 @@ struct SessionHandle { tasks: TaskManagerHandle, } +struct SessionSnapshot { + canonical_transcript: Vec, + config_options: Vec, +} + struct AttachedSession { session_id: wire::SessionId, config_options: Vec, @@ -720,6 +937,7 @@ impl Server { .collect(), connection, claim, + None, ) .await?; let AttachedSession { @@ -755,6 +973,25 @@ impl Server { )); } }; + // Validate before consulting the loaded actor: successful lookup emits + // its diagnostic identity, which a rejected resume must not change. + validate_resume_location( + self.runtime.root(), + &request.cwd.0, + !request.additional_directories.is_empty(), + )?; + if let Some(attached) = self.loaded_session(&request.session_id).await? { + let updates = if replay { + transcript_replay(&attached.session_id, &attached.canonical_transcript) + } else { + Vec::new() + }; + return Ok(( + wire::ResumeSessionResponse::new().config_options(attached.config_options), + updates, + attached.activation, + )); + } let claim = self .runtime .claim_session_load(&request.session_id.to_string())?; @@ -779,6 +1016,7 @@ impl Server { .collect(), connection, claim, + None, ) .await?; let updates = if replay { @@ -839,6 +1077,7 @@ impl Server { additional_directories: Vec, connection: V2ConnectionTo, mut claim: crate::runtime::SessionClaim, + forked: Option, ) -> Result { // Reject exhaustion before admission or any binding, driver, or actor effects. let token = self.registry.next_token()?; @@ -848,7 +1087,8 @@ impl Server { .map_err(|()| AcpRuntimeError::ClientClosed)?; let session_id = wire::SessionId::new(claim.id()); let cancellation = CancellationController::new(); - let sink = ResponseReplacementSink::new(ConnectionSink(connection)); + let injections = Arc::new(Mutex::new(InjectionWork::default())); + let sink = ResponseReplacementSink::new(ConnectionSink(connection, injections.clone())); let activity = native_activity(session_id.clone(), sink.clone()); let binding = AcpSessionBinding::new(session_id.clone(), SessionId::new(claim.id()), sink.clone()) @@ -871,7 +1111,20 @@ impl Server { cancellation: handle.cancellation_handle(), response_attempt_replacement: true, }; - let driver = self.runtime.start_acp_driver(context, &mut claim).await?; + // Admission starts before publication, not when the response/replay gate + // opens. Neither cancel nor close may become the new turn's baseline. + let initial_generation = forked + .as_ref() + .map(|_| handle.cancellation_handle().generation()); + handle.prepare_injection_turn(); + let activate_prompt = initial_generation.is_some(); + let driver = if let Some(forked) = forked { + self.runtime + .start_acp_branch_driver_with_initial(context, &mut claim, forked) + .await? + } else { + self.runtime.start_acp_driver(context, &mut claim).await? + }; let current = driver.adapter.selection().map_err(AcpRuntimeError::Loop)?; let reasoning = driver .adapter @@ -887,8 +1140,14 @@ impl Server { let structured_completion = driver.structured_completion; let mcp_events = self.runtime.subscribe_mcp(session_id.to_string()); let (tx, rx) = mpsc::channel(8); - let busy = Arc::new(AtomicBool::new(false)); + let busy = Arc::new(AtomicBool::new(activate_prompt)); let actor = SessionActor { + initial_generation, + admission_released: injections + .lock() + .expect("injection work poisoned") + .admission_released + .clone(), session_id: session_id.clone(), runtime: Arc::clone(&self.runtime), integration: Arc::clone(&self.integration), @@ -920,6 +1179,9 @@ impl Server { let _guard = guard; if activated.await.is_ok() { session_actor(actor).await; + } else { + // Abandoning response/replay retires the queued input as well. + actor.busy.store(false, Ordering::Release); } }); let interrupt_handle = handle.clone(); @@ -931,7 +1193,9 @@ impl Server { let weak = tx.downgrade(); let close_background_jobs = background_jobs.clone(); let close_tasks = tasks.clone(); + let close_handle = handle.clone(); let close = Arc::new(move || { + close_handle.close(); let close_background_jobs = close_background_jobs.clone(); let close_tasks = close_tasks.clone(); let weak = weak.clone(); @@ -957,6 +1221,7 @@ impl Server { completed: completion, session_id: session_id.clone(), session: SessionHandle { + injections, token, commands: tx, integration: handle, @@ -977,9 +1242,7 @@ impl Server { SessionPublicationError::Commit(error) => error, }); } - crate::events::emit(&crate::events::RuntimeEvent::SessionStarted { - session_id: session_id.to_string(), - }); + establish_diagnostic_route(&session_id); drop(actor_task); Ok(AttachedSession { session_id, @@ -989,6 +1252,191 @@ impl Server { }) } + // Replaying a session already owned by this connection is read-only. It + // must not reacquire its disk lock, replace its adapter, or restart a turn. + async fn loaded_session( + &self, + session_id: &wire::SessionId, + ) -> Result, AcpRuntimeError> { + let sender = self + .sessions + .lock() + .expect("ACP v2 session map poisoned") + .get(session_id) + .map(|session| session.commands.clone()); + let Some(sender) = sender else { + return Ok(None); + }; + let (reply, response) = oneshot::channel(); + sender + .send(Command::Snapshot { reply }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + let SessionSnapshot { + canonical_transcript, + config_options, + } = response + .await + .map_err(|_| AcpRuntimeError::ClientClosed)??; + // Reattachment can switch the client's diagnostic route without creating + // an actor. Write the source marker before returning the load response. + establish_diagnostic_route(session_id); + let (activation, _already_active) = oneshot::channel(); + Ok(Some(AttachedSession { + session_id: session_id.clone(), + config_options, + canonical_transcript, + activation, + })) + } + + fn branch_admission( + &self, + session_id: &wire::SessionId, + ) -> Result<(mpsc::Sender, BranchAdmission), AcpRuntimeError> { + let sessions = self.sessions.lock().expect("ACP v2 session map poisoned"); + let session = sessions + .get(session_id) + .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string()))?; + let admission = BranchAdmission::claim(session.busy.clone(), session.injections.clone())?; + Ok((session.commands.clone(), admission)) + } + + async fn list_prompt_branches( + &self, + request: ListPromptBranchesRequest, + ) -> Result { + let (sender, _, _) = self.prompt_route(&request.session_id)?; + let (reply, response) = oneshot::channel(); + sender + .send(Command::ListPromptBranches { reply }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + response.await.map_err(|_| AcpRuntimeError::ClientClosed)? + } + + async fn prepare_prompt_branch( + &self, + request: PreparePromptBranchRequest, + ) -> Result { + let (sender, admission) = self.branch_admission(&request.session_id)?; + let (reply, response) = oneshot::channel(); + sender + .send(Command::PreparePromptBranch { + address: request.address, + admission, + reply, + }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + response.await.map_err(|_| AcpRuntimeError::ClientClosed)? + } + + async fn submit_prompt_branch( + self: &Arc, + request: SubmitPromptBranchRequest, + connection: V2ConnectionTo, + ) -> Result<(SubmitPromptBranchResponse, Option), AcpRuntimeError> { + // Idempotency is global, not tied to a live actor or its busy flag. A + // retry after restart/close/source advancement must find the disk child. + let _submission = BRANCH_SUBMISSIONS.lock().await; + if let Some(committed) = prompt_branches::lookup_committed( + self.runtime.root(), + &request.session_id.to_string(), + &request.checkout_token, + &request.text, + ) + .map_err(AcpRuntimeError::Loop)? + { + let child_id = wire::SessionId::new(committed.session_id.clone()); + if let Some(attached) = self + .loaded_session(&child_id) + .await + .map_err(|error| branch_child_error(&committed.session_id, error))? + { + let response = SubmitPromptBranchResponse { + session_id: attached.session_id.clone(), + config_options: attached.config_options.clone(), + }; + return Ok((response, Some(attached))); + } + let claim = self + .runtime + .claim_session_load(&committed.session_id) + .map_err(|error| branch_child_error(&committed.session_id, error))?; + let attached = self + .attach_session( + self.runtime.root().to_owned(), + Vec::new(), + connection, + claim, + None, + ) + .await + .map_err(|error| branch_child_error(&committed.session_id, error))?; + let response = SubmitPromptBranchResponse { + session_id: attached.session_id.clone(), + config_options: attached.config_options.clone(), + }; + // A durable retry may reattach, but never starts another generation. + return Ok((response, Some(attached))); + } + let (sender, admission) = self.branch_admission(&request.session_id)?; + let (reply, response) = oneshot::channel(); + sender + .send(Command::ReservePromptBranch { + checkout_token: request.checkout_token, + admission, + reply, + }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + let (checkout, release) = response + .await + .map_err(|_| AcpRuntimeError::ClientClosed)??; + let forked = checkout + .fork(&request.text) + .map_err(AcpRuntimeError::Loop)?; + let claim = self.runtime.claim_session_fork()?; + let child_id = claim.id().to_string(); + // Keep the actor reservation alive through the durable commit and route + // publication. Dropping release also unblocks the actor on any failure. + let attached = self + .attach_session( + self.runtime.root().to_owned(), + Vec::new(), + connection, + claim, + Some(forked), + ) + .await; + drop(release); + let attached = attached.map_err(|error| branch_child_error(&child_id, error))?; + let response = SubmitPromptBranchResponse { + session_id: attached.session_id.clone(), + config_options: attached.config_options.clone(), + }; + Ok((response, Some(attached))) + } + + fn injection_admission( + &self, + session_id: &wire::SessionId, + ) -> Result { + let sessions = self.sessions.lock().expect("ACP v2 session map poisoned"); + let session = sessions + .get(session_id) + .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string()))?; + let mut work = session.injections.lock().expect("injection work poisoned"); + if work.branch_reserved { + return Err(AcpRuntimeError::Loop( + "source session is reserved for prompt checkout".into(), + )); + } + work.admitting += 1; + Ok(InjectionAdmission(session.injections.clone())) + } + async fn prepare_prompt( &self, request: wire::PromptRequest, @@ -1095,10 +1543,11 @@ impl Server { ) }); if let Some((handle, background_jobs, tasks, structured_completion)) = session { + // Interrupt admission before waiting for asynchronous task cleanup. + handle.interrupt(); if structured_completion { super::cancel_background_jobs(&tasks, &background_jobs).await; } - handle.interrupt(); } Ok(()) } @@ -1113,17 +1562,19 @@ impl Server { .map_err(|_| AcpRuntimeError::ClientClosed)? .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - super::cancel_background_jobs(&session.tasks, &session.background_jobs).await; session.integration.close(); + super::cancel_background_jobs(&session.tasks, &session.background_jobs).await; let (reply, acknowledged) = oneshot::channel(); - session + // A cancelled, unactivated child may retire its actor before this + // command arrives. A dropped receiver is already a completed close. + if session .commands .send(Command::Close { reply }) .await - .map_err(|_| AcpRuntimeError::ClientClosed)?; - acknowledged - .await - .map_err(|_| AcpRuntimeError::ClientClosed)?; + .is_ok() + { + let _ = acknowledged.await; + } self.registry.remove(session.token); Ok(wire::CloseSessionResponse::new()) } @@ -1163,14 +1614,62 @@ impl Server { } } -struct SessionActor { +async fn hold_branch_reservation( + checkout: T, + admission: BranchAdmission, + reply: oneshot::Sender), AcpRuntimeError>>, +) { + let (release, released) = oneshot::channel(); + if reply.send(Ok((checkout, release))).is_ok() { + // No select here: config, MCP and autonomous work must remain queued + // until the child is committed. Cancellation also releases the actor. + let _ = released.await; + } + drop(admission); +} + +fn branch_child_error(child_id: &str, error: AcpRuntimeError) -> AcpRuntimeError { + AcpRuntimeError::Loop(format!( + "prompt checkout child {child_id}: {error}; retry the same checkout to recover a committed child" + )) +} + +async fn settled_branch_source( + driver: &LoopDriver, + tasks: &TaskManagerHandle, + jobs: &BackgroundJobs, + mcp: &crate::tools::mcp::McpSubscription, +) -> Result<(), AcpRuntimeError> { + let before = jobs.activity(); + let running = !tasks.list_running().await.is_empty(); + let after = jobs.activity(); + if !driver.snapshot().pending_input.is_empty() + || running + || before.active + || after.active + || before.unacknowledged_terminals + || after.unacknowledged_terminals + || before.generation != after.generation + || mcp.has_pending() + || driver.wait_for_loop_update().now_or_never().is_some() + { + return Err(AcpRuntimeError::Loop( + "source session has unsettled work".into(), + )); + } + Ok(()) +} + +struct SessionActor { + initial_generation: Option, + admission_released: Arc, session_id: wire::SessionId, runtime: Arc, integration: Arc, handle: AcpSessionHandle, busy: Arc, binding: BindingGuard, - sink: ResponseReplacementSink, + sink: ResponseReplacementSink, activity: SessionActivity, driver: LoopDriver, tasks: TaskManagerHandle, @@ -1183,8 +1682,69 @@ struct SessionActor { mcp_events: crate::tools::mcp::McpSubscription, } -async fn session_actor(actor: SessionActor) { +// Content equality does not establish uninterrupted authority: admitted work +// (including compaction) can restore identical bytes. Read-only actor commands +// never advance this shared transcript/configuration revision. +fn advance_checkout_revision(revision: &mut u64) { + *revision = revision + .checked_add(1) + .expect("checkout revision exhausted"); +} + +#[allow(clippy::too_many_arguments)] +async fn run_initial_branch_turn( + session_id: &wire::SessionId, + integration: &AcpIntegration, + handle: &AcpSessionHandle, + busy: &AtomicBool, + mut driver: LoopDriver, + sink: &ResponseReplacementSink, + generation: u64, + structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, + activity: &SessionActivity, +) -> Result>, AcpRuntimeError> { + // Runtime already committed this user item and queued it exactly once. + // Admission and its cancellation baseline precede the publication gate. + integration.finish_prompt(session_id); + if !handle.cancellation_handle().is_cancelled_since(generation) { + handle.start_injection_turn(); + } + let result = run_active_turn( + session_id, + integration, + handle, + &mut driver, + sink, + generation, + structured, + activity, + ExecutionOrigin::Autonomous, + ) + .await; + handle.stop_injection_turn(); + busy.store(false, Ordering::Release); + if !driver.is_available() || !driver.snapshot().pending_input.is_empty() { + result?; + // retire_interrupted_turn preserves queued input. There is no loop API + // to clear unstarted input, so retire this attachment instead. Dropping + // the driver prevents a later MCP/task/injection wake from executing it; + // the committed child remains discoverable and reloads passively. + return Ok(None); + } + // As with ordinary prompts, an execution error is already reported and + // settled. Once input was consumed, retain the driver for the next prompt. + if let Err(error) = result { + eprintln!("ACP v2 branch turn failed for {session_id}: {error}"); + } + Ok(Some(driver)) +} + +async fn session_actor( + actor: SessionActor, +) { let SessionActor { + initial_generation, + admission_released, session_id, runtime, integration, @@ -1205,21 +1765,111 @@ async fn session_actor(actor: SessionActor) } = actor; let mut binding = Some(binding); let mut model_switch = model_switch::Guard::default(); + let mut checkouts = PromptCheckouts::default(); + let mut checkout_revision = 0u64; + // Event selection is not admission: a branch may claim busy before its + // command reaches this actor. Retain the wake until a drive is admitted. + let mut autonomous_pending = false; + if let Some(generation) = initial_generation { + advance_checkout_revision(&mut checkout_revision); + let initial = run_initial_branch_turn( + &session_id, + &integration, + &handle, + &busy, + driver, + &sink, + generation, + structured_completion.then_some((&tasks, &background_jobs)), + &activity, + ) + .await; + match initial { + Ok(Some(active)) => driver = active, + stopped => { + if let Err(error) = stopped { + eprintln!("ACP v2 branch turn failed for {session_id}: {error}"); + } + super::cancel_background_jobs(&tasks, &background_jobs).await; + drop(binding.take()); + commands.close(); + while let Some(command) = commands.recv().await { + if let Command::Close { reply } = command { + let _ = reply.send(()); + } + } + return; + } + } + } loop { tokio::select! { biased; command = commands.recv() => match command { - Some(Command::Prompt(command)) => { - let result = prepare_prompt( - &session_id, - &runtime, - &integration, - &handle, - &mut skill_catalog, - &mut driver, - command, - &sink, - &tasks, + Some(Command::Snapshot { reply }) => { + let result = (|| { + let selection = adapter.selection().map_err(AcpRuntimeError::Loop)?; + let reasoning = adapter.reasoning_effort().map_err(AcpRuntimeError::Loop)?; + Ok(SessionSnapshot { canonical_transcript: driver.snapshot().transcript, config_options: v2_config_options(&selection, reasoning, &catalog) }) + })(); + let _ = reply.send(result); + } + Some(Command::ListPromptBranches { reply }) => { + let result = (|| { + let selection = adapter.selection().map_err(AcpRuntimeError::Loop)?; + let reasoning = adapter.reasoning_effort().map_err(AcpRuntimeError::Loop)?; + checkouts.list(runtime.root(), &session_id.to_string(), &driver.snapshot().transcript, + checkout_revision, &selection, reasoning).map_err(AcpRuntimeError::Loop) + })(); + let _ = reply.send(result); + } + Some(Command::PreparePromptBranch { address, admission, reply }) => { + let result = async { + settled_branch_source(&driver, &tasks, &background_jobs, &mcp_events).await?; + let selection = adapter.selection().map_err(AcpRuntimeError::Loop)?; + let reasoning = adapter.reasoning_effort().map_err(AcpRuntimeError::Loop)?; + let checkout = checkouts.prepare(&address, &driver.snapshot().transcript, + checkout_revision, &selection, reasoning).map_err(AcpRuntimeError::Loop)?; + Ok(PreparePromptBranchResponse { + checkout_token: checkout.token, + original_text: checkout.original_text, + prefix: transcript_replay(&session_id, &checkout.prefix).into_iter().map(|update| update.update).collect(), + config_options: v2_config_options(&checkout.selection, checkout.reasoning, &catalog), + }) + }.await; + drop(admission); + let _ = reply.send(result); + } + Some(Command::ReservePromptBranch { checkout_token, admission, reply }) => { + let result = async { + settled_branch_source(&driver, &tasks, &background_jobs, &mcp_events).await?; + let selection = adapter.selection().map_err(AcpRuntimeError::Loop)?; + let reasoning = adapter.reasoning_effort().map_err(AcpRuntimeError::Loop)?; + checkouts.checkout(&checkout_token, &driver.snapshot().transcript, + checkout_revision, &selection, reasoning).map_err(AcpRuntimeError::Loop) + }.await; + match result { + Ok(checkout) => { + hold_branch_reservation(checkout, admission, reply).await; + continue; + } + Err(error) => { let _ = reply.send(Err(error)); } + } + drop(admission); + } + Some(Command::Prompt(command)) => { + let generation = command.cancellation_generation; + advance_checkout_revision(&mut checkout_revision); + let result = prepare_prompt( + &session_id, + &runtime, + &integration, + &handle, + &mut skill_catalog, + &mut driver, + command, + &sink, + &tasks, &background_jobs, structured_completion, &activity,) @@ -1229,8 +1879,18 @@ async fn session_actor(actor: SessionActor) if let Err(error) = result { eprintln!("ACP v2 prompt failed for {session_id}: {error}"); } + if !driver.is_available() || (handle.cancellation_handle().is_cancelled_since(generation) + && !driver.snapshot().pending_input.is_empty()) + { + // The same pre-step cancellation race applies to a + // queued ordinary prompt, not just branch activation. + let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); + super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; + break; + } } Some(Command::SetConfig { request, reply, cancellation_generation }) => { + let mut compaction_started = false; let result = async { if handle.cancellation_handle().is_cancelled_since(cancellation_generation) { return Err(model_switch::error("model change cancelled")); @@ -1247,6 +1907,8 @@ async fn session_actor(actor: SessionActor) claim_prompt(&busy).map_err(sdk_error)?; handle.prepare_injection_turn(); handle.start_injection_turn(); + advance_checkout_revision(&mut checkout_revision); + compaction_started = true; let compacted = compact_for_switch( &session_id, &integration, &handle, &mut driver, &sink, cancellation_generation, &activity, @@ -1261,7 +1923,19 @@ async fn session_actor(actor: SessionActor) } set_v2_config(&adapter, &catalog, request).map_err(sdk_error) }.await; + if result.is_ok() { advance_checkout_revision(&mut checkout_revision); } let _ = reply.send(result); + if !driver.is_available() || (compaction_started + && handle.cancellation_handle().is_cancelled_since(cancellation_generation) + && !driver.snapshot().pending_input.is_empty()) + { + // Cancellation before the first step leaves the synthetic + // /compact input queued. As with an interrupted prompt, + // retire this attachment so no later prompt or wake runs it. + let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); + super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; + break; + } } Some(Command::Close { reply }) => { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); @@ -1276,45 +1950,47 @@ async fn session_actor(actor: SessionActor) break; } }, + _ = admission_released.notified(), if autonomous_pending => {} + _ = std::future::ready(()), if autonomous_pending && !busy.load(Ordering::Acquire) => { + let generation = handle.cancellation_handle().generation(); + let result = drive_autonomous( + &session_id, &integration, &handle, &busy, + &mut driver, &sink, &activity, + ) + .instrument(crate::telemetry::error_spans::operation("acp_autonomous")) + .await; + autonomous_pending = matches!(result, Ok(false)); + if let Err(error) = result { + eprintln!("ACP v2 autonomous turn failed for {session_id}: {error}"); + } + if !driver.is_available() || (!autonomous_pending + && handle.cancellation_handle().is_cancelled_since(generation) + && !driver.snapshot().pending_input.is_empty()) + { + let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); + super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; + break; + } + } event = mcp_events.recv() => { if let Some(event) = event { - let result = async { - match driver.submit_input(vec![Item::notification(event.message)]) { - Ok(()) => drive_autonomous( - &session_id, - &integration, - &handle, - &busy, - &mut driver, - &sink, - &activity, - ).await, - Err(error) => Err(map_loop_error(&session_id, &error)), + advance_checkout_revision(&mut checkout_revision); + match driver.submit_input(vec![Item::notification(event.message)]) { + Ok(()) => { + autonomous_pending = true; + } + Err(error) => { + eprintln!("ACP v2 autonomous turn failed for {session_id}: {}", map_loop_error(&session_id, &error)); } - } - .instrument(crate::telemetry::error_spans::operation("acp_autonomous")) - .await; - if let Err(error) = result { - eprintln!("ACP v2 autonomous turn failed for {session_id}: {error}"); } } } event = tasks.next_event() => match event { Some(TaskEvent::Completed(snapshot, _)) => { background_jobs.acknowledge_terminal(&snapshot.call_id); - if snapshot.kind == agentkit_task_manager::TaskKind::Background - && let Err(error) = drive_autonomous( - &session_id, - &integration, - &handle, - &busy, - &mut driver, - &sink, - &activity,) - .instrument(crate::telemetry::error_spans::operation("acp_autonomous")) - .await - { - eprintln!("ACP v2 autonomous turn failed for {session_id}: {error}"); + if snapshot.kind == agentkit_task_manager::TaskKind::Background { + advance_checkout_revision(&mut checkout_revision); + autonomous_pending = true; } } Some(TaskEvent::Cancelled(snapshot) | TaskEvent::Failed(snapshot, _)) => { @@ -1422,25 +2098,80 @@ async fn drive_prompt( handle: &AcpSessionHandle, cancellation_generation: u64, structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, + acknowledgements: &SteeringAcknowledgements, ) -> Result where S: ModelSession + Send + 'static, { + let mut acknowledged_pending = Vec::new(); let result = drive_prompt_inner( session_id, driver, handle, cancellation_generation, structured, + acknowledgements, + &mut acknowledged_pending, ) .await; if matches!(result, Ok(FinishReason::Cancelled)) { - // A cooperative interrupt is still a live logical turn. Retire it - // without another `next`, which could execute cancelled model work. - driver - .retire_interrupted_turn() - .await - .map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; + if !acknowledged_pending.is_empty() && !driver.snapshot().pending_input.is_empty() { + // A successful injection boundary can stop after acknowledging one + // steer while awaiting another response's activation. Those items + // are accepted input, not the unstarted original prompt. Commit + // them through the real transcript observer with execution fenced. + driver + .settle_delivered_input() + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; + } else { + // Unstarted original/synthetic input must not be drained here. The + // actor retires that attachment so no later wake executes it. + driver + .retire_interrupted_turn() + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; + } + } + result +} + +// A partial boundary can contain acknowledged A and submitted but unacknowledged +// B. Only the exact ordered prefix witnessed by successful notifications is +// durable. A failed notification retires the attachment, never retries B. +async fn injection_boundary( + driver: &mut LoopDriver, + handle: &AcpSessionHandle, + terminal: bool, + acknowledgements: &SteeringAcknowledgements, + acknowledged_pending: &mut Vec, +) -> Result { + if driver.snapshot().pending_input != *acknowledged_pending { + driver.make_unavailable(); + return Err(AcpRuntimeError::Loop( + "unclassified input at injection boundary".into(), + )); + } + let boundary = acknowledgements + .begin() + .inspect_err(|_| driver.make_unavailable())?; + let result = handle.handle_injection_boundary(driver, terminal).await; + acknowledged_pending.extend( + boundary + .finish() + .inspect_err(|_| driver.make_unavailable())?, + ); + if result.is_ok() && driver.snapshot().pending_input != *acknowledged_pending { + driver.make_unavailable(); + return Err(AcpRuntimeError::Loop( + "injection acknowledgement did not match queued items".into(), + )); + } + if result.is_err() { + let accepted = std::mem::take(acknowledged_pending); + let settled = driver.settle_acknowledged_input(&accepted).await; + driver.make_unavailable(); + settled.map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; } result } @@ -1451,11 +2182,22 @@ async fn drive_prompt_inner( handle: &AcpSessionHandle, cancellation_generation: u64, structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, + acknowledgements: &SteeringAcknowledgements, + acknowledged_pending: &mut Vec, ) -> Result where S: ModelSession + Send + 'static, { loop { + // Check before *every* driver step: next() can checkpoint a fresh + // generation and dispatch queued input before its first await returns. + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { + handle.stop_injection_turn(); + return Ok(FinishReason::Cancelled); + } let step = match driver.next().await { Ok(step) => step, Err(error) => { @@ -1469,6 +2211,9 @@ where return loop_error_stop_reason(session_id, &error); } }; + if driver.snapshot().pending_input.is_empty() { + acknowledged_pending.clear(); + } if handle .cancellation_handle() .is_cancelled_since(cancellation_generation) @@ -1495,7 +2240,15 @@ where { continue; } - match handle.handle_injection_boundary(driver, true).await { + match injection_boundary( + driver, + handle, + true, + acknowledgements, + acknowledged_pending, + ) + .await + { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; } @@ -1523,7 +2276,15 @@ where { continue; } - match handle.handle_injection_boundary(driver, true).await { + match injection_boundary( + driver, + handle, + true, + acknowledgements, + acknowledged_pending, + ) + .await + { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; } @@ -1546,7 +2307,15 @@ where } } LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => { - match handle.handle_injection_boundary(driver, false).await { + match injection_boundary( + driver, + handle, + false, + acknowledgements, + acknowledged_pending, + ) + .await + { Ok(AcpInjectionBoundary::Stopped) => { return Ok(FinishReason::Cancelled); } @@ -1637,6 +2406,10 @@ async fn run_active_turn( activity: &SessionActivity, origin: ExecutionOrigin, ) -> Result { + // A failed child submission may leave stderr routed to the child even + // after the client returns to its source without sending session/resume. + // Re-establish identity before activity, model, or cleanup diagnostics. + establish_diagnostic_route(session_id); activity .execute( origin, @@ -1647,6 +2420,7 @@ async fn run_active_turn( handle, cancellation_generation, structured, + &sink.acknowledgements, ) .await; let outcome = super::activity::ExecutionOutcome::new( @@ -1671,6 +2445,7 @@ async fn run_active_turn( .await } +// False means admission was denied; the actor must retain the selected wake. async fn drive_autonomous( session_id: &wire::SessionId, integration: &AcpIntegration, @@ -1679,13 +2454,13 @@ async fn drive_autonomous( driver: &mut LoopDriver, sink: &ResponseReplacementSink, activity: &SessionActivity, -) -> Result<(), AcpRuntimeError> { +) -> Result { + let cancellation_generation = handle.cancellation_handle().generation(); if claim_prompt(busy).is_err() { - return Ok(()); + return Ok(false); } handle.prepare_injection_turn(); integration.finish_prompt(session_id); - let cancellation_generation = handle.cancellation_handle().generation(); handle.start_injection_turn(); let result = run_active_turn( session_id, @@ -1702,7 +2477,7 @@ async fn drive_autonomous( integration.finish_prompt(session_id); handle.stop_injection_turn(); busy.store(false, Ordering::Release); - result.map(|_| ()) + result.map(|_| true) } fn error_diagnostic_notification( @@ -2133,6 +2908,92 @@ pub(crate) fn component( }, agent_client_protocol::on_receive_request!(), ) + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: ListPromptBranchesRequest, responder, cx| { + let state = Arc::clone(&state); + cx.spawn(async move { + responder.respond_with_result( + state.list_prompt_branches(request).await.map_err(sdk_error), + ) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: PreparePromptBranchRequest, responder, cx| { + let state = Arc::clone(&state); + cx.spawn(async move { + responder.respond_with_result( + state + .prepare_prompt_branch(request) + .await + .map_err(sdk_error), + ) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: SubmitPromptBranchRequest, + responder: Responder, + cx| { + let state = Arc::clone(&state); + let connection = cx.clone(); + cx.spawn(async move { + match state + .submit_prompt_branch(request, connection.clone()) + .await + { + Ok((response, attached)) => { + let child_id = response.session_id.clone(); + let postcommit = |error: agent_client_protocol::Error| { + agent_client_protocol::util::internal_error(format!( + "prompt checkout child {child_id}: {error}" + )) + }; + // The response establishes the child route before + // replay; execution is gated behind both phases. + responder + .respond_tracked(response) + .map_err(&postcommit)? + .await + .map_err(&postcommit)?; + if let Some(attached) = attached { + for update in transcript_replay( + &attached.session_id, + &attached.canonical_transcript, + ) { + connection + .send_notification(update) + .map_err(&postcommit)?; + } + connection + .send_notification(available_commands_update( + child_id.clone(), + )) + .map_err(&postcommit)?; + let _ = attached.activation.send(()); + } + Ok(()) + } + Err(error) => responder.respond_with_error(sdk_error(error)), + } + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) .on_receive_request( { let state = Arc::clone(&state); @@ -2168,13 +3029,40 @@ pub(crate) fn component( ) .on_receive_request( { - let integration = Arc::clone(&state.integration); + let state = Arc::clone(&state); async move |request: wire::InjectSessionRequest, responder: Responder, cx| { - let integration = Arc::clone(&integration); + let admission = match state.injection_admission(&request.session_id) { + Ok(admission) => admission, + Err(error) => return responder.respond_with_error(sdk_error(error)), + }; + let integration = Arc::clone(&state.integration); cx.spawn(async move { - integration.handle_inject_request(request, responder).await + let Some(reserved) = integration + .reserve_inject_request(request, responder) + .await? + else { + return Ok(()); + }; + let id = reserved.response().message_id; + admission + .0 + .lock() + .expect("injection work poisoned") + .pending + .insert(id.clone()); + let mut pending = TrackedInjection { + work: admission.0.clone(), + id, + retained: false, + }; + if let Some(acceptance) = reserved.respond_tracked()? { + acceptance.activate_after_response().await?; + pending.retained = true; + } + drop(admission); + Ok(()) })?; Ok(()) } @@ -2183,9 +3071,26 @@ pub(crate) fn component( ) .on_receive_request( { - let integration = Arc::clone(&state.integration); + let state = Arc::clone(&state); async move |request: wire::RevokeInjectSessionRequest, responder, _cx| { - responder.respond_with_result(integration.revoke_inject(request).await) + let tracker = state + .sessions + .lock() + .expect("ACP v2 session map poisoned") + .get(&request.session_id) + .map(|session| session.injections.clone()); + let id = request.message_id.clone(); + let result = state.integration.revoke_inject(request).await; + if result.is_ok() + && let Some(tracker) = tracker + { + tracker + .lock() + .expect("injection work poisoned") + .pending + .remove(&id); + } + responder.respond_with_result(result) } }, agent_client_protocol::on_receive_request!(), @@ -2253,48 +3158,1383 @@ pub(crate) fn component( Ok(Handled::Yes) } }, - agent_client_protocol::on_receive_notification!(), - ) - .on_receive_request( - { - let state = Arc::clone(&state); - async move |request: wire::CloseSessionRequest, responder, cx| { - let state = Arc::clone(&state); - cx.spawn(async move { - responder.respond_with_result(state.close(request).await.map_err(sdk_error)) - })?; - Ok(()) + agent_client_protocol::on_receive_notification!(), + ) + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: wire::CloseSessionRequest, responder, cx| { + let state = Arc::clone(&state); + cx.spawn(async move { + responder.respond_with_result(state.close(request).await.map_err(sdk_error)) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ); + Ok(agent) +} + +#[cfg(test)] +mod tests { + mod migration_resume; + mod review_tests; + + use std::{ + collections::VecDeque, + sync::{atomic::AtomicUsize, mpsc as std_mpsc}, + }; + + use serde_json::json; + + use agent_client_protocol::schema::MaybeUndefined; + use agentkit_core::{DataRef, MetadataMap, Modality, TurnCancellation}; + use agentkit_loop::{ + Agent, ModelAdapter, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, + TurnRequest, TurnResult, + }; + use agentkit_task_manager::{AsyncTaskManager, RoutingDecision, TaskManager}; + use agentkit_tools_core::{ToolAnnotations, ToolName, ToolRegistry, ToolSpec}; + use tokio::{ + sync::Notify, + time::{Duration, timeout}, + }; + + use super::*; + use crate::protocols::acp::tests::{BlockingTool, ScriptAdapter}; + + #[tokio::test] + async fn branch_admitted_work_invalidates_checkout_even_when_bytes_do_not_change() { + let root = tempfile::tempdir().unwrap(); + let source_id = crate::session::new_id(); + let source = crate::session::open( + root.path(), + &source_id, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "original prompt"), + Item::text(ItemKind::Assistant, "original answer"), + ], + ) + .unwrap(); + let original = source.transcript.clone(); + drop(source); + let selection = + crate::provider::ModelSelection::new(ProviderKind::OpenRouter, "test-model"); + let mut checkouts = PromptCheckouts::default(); + let mut revision = 0; + let listed = checkouts + .list( + root.path(), + &source_id, + &original, + revision, + &selection, + None, + ) + .unwrap(); + let checkout = checkouts + .prepare( + &listed.boundaries[0].address, + &original, + revision, + &selection, + None, + ) + .unwrap(); + let integration = AcpIntegration::default(); + let sink = ResponseReplacementSink::new(RecordingSink::default()); + let session_id = wire::SessionId::new(source_id.clone()); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(source_id.clone()), + sink.clone(), + )) + .unwrap(); + let turns = Arc::new(AtomicU64::new(0)); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .transcript(original.clone()) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(source_id.clone())).without_cache()) + .await + .unwrap(); + let mut driver = input_settlement.wrap(driver); + // Re-listing and reading an identical snapshot leave the token valid. + checkouts + .list( + root.path(), + &source_id, + &driver.snapshot().transcript, + revision, + &selection, + None, + ) + .unwrap(); + checkouts + .checkout( + &checkout.token, + &driver.snapshot().transcript, + revision, + &selection, + None, + ) + .unwrap(); + // A background wake can settle without changing canonical bytes. Its + // admission still invalidates authority, just like a compaction round + // that restores the same contents. + advance_checkout_revision(&mut revision); + drive_autonomous( + &session_id, + &integration, + &handle, + &AtomicBool::new(false), + &mut driver, + &sink, + &activity, + ) + .await + .unwrap(); + assert_eq!(driver.snapshot().transcript, original); + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert!( + checkouts + .checkout( + &checkout.token, + &driver.snapshot().transcript, + revision, + &selection, + None + ) + .is_err() + ); + } + + #[tokio::test] + async fn branch_initial_activation_drives_committed_pending_prompt_exactly_once() { + let integration = AcpIntegration::default(); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let session_id = wire::SessionId::new("initial-branch"); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new("initial-branch"), + sink.clone(), + )) + .unwrap(); + let turns = Arc::new(AtomicU64::new(0)); + let prompt = Item::text(ItemKind::User, "already committed edited prompt"); + // Match runtime's fresh-branch bootstrap: disk/canonical replay already + // has the prompt, while the loop receives it once through pending input. + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .observer(ResponseReplacementObserver::new( + integration.clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .transcript(vec![Item::text(ItemKind::System, "system")]) + .input(vec![prompt.clone()]) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new("initial-branch")).without_cache()) + .await + .unwrap(); + let driver = input_settlement.wrap(driver); + assert_eq!(driver.snapshot().pending_input.len(), 1); + let busy = Arc::new(AtomicBool::new(true)); + let actor_busy = busy.clone(); + handle.prepare_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let (activation, activated) = oneshot::channel(); + let actor = tokio::spawn(async move { + activated.await.unwrap(); + let driver = run_initial_branch_turn( + &session_id, + &integration, + &handle, + &actor_busy, + driver, + &sink, + generation, + None, + &activity, + ) + .await + .unwrap() + .unwrap(); + driver.snapshot() + }); + tokio::task::yield_now().await; + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert!(recording.updates.lock().unwrap().is_empty()); + assert!(busy.load(Ordering::Acquire)); + // The production actor opens this same gate only after response/replay. + activation.send(()).unwrap(); + let snapshot = actor.await.unwrap(); + assert_eq!(turns.load(Ordering::Relaxed), 1); + assert!(snapshot.pending_input.is_empty()); + assert_eq!( + snapshot + .transcript + .iter() + .filter(|item| item.kind == ItemKind::User) + .count(), + 1 + ); + assert_eq!( + snapshot + .transcript + .iter() + .find(|item| item.kind == ItemKind::User) + .unwrap() + .parts, + prompt.parts + ); + assert_eq!( + snapshot + .transcript + .iter() + .filter(|item| item.kind == ItemKind::Assistant) + .count(), + 1 + ); + assert!(!busy.load(Ordering::Acquire)); + assert_running_then_idle( + &recording.updates.lock().unwrap(), + wire::StopReason::EndTurn, + ); + assert!( + !recording + .updates + .lock() + .unwrap() + .iter() + .any(|update| matches!(update.update, wire::SessionUpdate::UserMessage(_))) + ); + } + + #[tokio::test] + async fn branch_first_turn_errors_allow_next_prompt_on_same_actor() { + async fn snapshot(commands: &mpsc::Sender) -> SessionSnapshot { + let (reply, response) = oneshot::channel(); + commands.send(Command::Snapshot { reply }).await.unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .expect("execution errors must retain the loaded child") + .unwrap() + } + + for outcome in [ + TestOutcome::ProviderErrorThenContent, + TestOutcome::FinishErrorThenContent, + ] { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let child_id = crate::session::new_id(); + let session_id = wire::SessionId::new(child_id.clone()); + let selection = + crate::provider::ModelSelection::new(ProviderKind::OpenAiSubscription, "gpt-5.4"); + let prefix = vec![Item::text(ItemKind::System, "system")]; + let initial = crate::session::branch::prepare( + prefix.clone(), + "source".into(), + crate::session::branch::Boundary::new(0, &prefix).unwrap(), + "failed-first-turn".into(), + crate::session::branch::SubmittedRequest { + id: "edited-request".into(), + selection: crate::session::branch::CapturedSelection::new(&selection, None), + }, + Item::text(ItemKind::User, "committed edited prompt"), + ) + .unwrap(); + // Successful submit committed this exact prompt before execution. + // Runtime places it in the loop's pending input, not its prefix. + let child = + crate::session::open_uncommitted(root.path(), &child_id, false, initial).unwrap(); + crate::session::branch::commit(&child.observer, &child.transcript).unwrap(); + let mut prefix = child.transcript.clone(); + let prompt = prefix.pop().unwrap(); + drop(child); + + let integration = Arc::new(AcpIntegration::default()); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(child_id.clone()), + sink.clone(), + )) + .unwrap(); + handle.prepare_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let turns = Arc::new(AtomicU64::new(0)); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome, + turns: turns.clone(), + interrupt: None, + }) + .observer(ResponseReplacementObserver::new( + integration.as_ref().clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .task_manager(manager) + .transcript(prefix) + .input(vec![prompt.clone()]) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(child_id.clone())).without_cache()) + .await + .unwrap(); + let driver = input_settlement.wrap(driver); + let busy = Arc::new(AtomicBool::new(true)); + let (commands, receiver) = mpsc::channel(8); + let mcp_events = runtime.subscribe_mcp(child_id.clone()); + let actor = tokio::spawn(session_actor(SessionActor { + initial_generation: Some(generation), + admission_released: Arc::new(Notify::new()), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(), + catalog: vec![], + commands: receiver, + mcp_events, + })); + // Snapshot is a serialized actor barrier, not a scheduler delay. + let failed = snapshot(&commands).await; + assert_eq!(turns.load(Ordering::Relaxed), 1); + assert!(!busy.load(Ordering::Acquire)); + assert_eq!(failed.canonical_transcript.last(), Some(&prompt)); + assert_running_then_idle(&recording.updates.lock().unwrap(), error_stop_reason()); + + let (reply, response) = oneshot::channel(); + commands + .send(Command::SetConfig { + request: wire::SetSessionConfigOptionRequest::new( + session_id.clone(), + "reasoning_effort", + "high", + ), + reply, + cancellation_generation: handle.cancellation_handle().generation(), + }) + .await + .unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + turns.load(Ordering::Relaxed), + 1, + "config must not execute a model turn" + ); + + // Ordinary admission on this exact actor and binding: no load, + // switch, or replacement driver. A new fake session would fail again. + let generation = handle.cancellation_handle().generation(); + claim_prompt(&busy).unwrap(); + handle.prepare_injection_turn(); + let (reply, response) = oneshot::channel(); + commands + .send(Command::Prompt(PromptCommand { + request: wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "try again", + ))], + ), + cancellation_generation: generation, + reply, + })) + .await + .unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .unwrap() + .send(()) + .unwrap(); + let recovered = snapshot(&commands).await; + assert_eq!(turns.load(Ordering::Relaxed), 2); + assert!(!busy.load(Ordering::Acquire)); + assert_eq!( + recovered + .canonical_transcript + .iter() + .filter(|item| item.kind == ItemKind::User) + .count(), + 2 + ); + assert_eq!( + recovered + .canonical_transcript + .iter() + .filter(|item| item.kind == ItemKind::Assistant) + .count(), + 1 + ); + assert!( + matches!(recording.updates.lock().unwrap().last().map(|update| &update.update), + Some(wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle))) + if idle.stop_reason.as_ref() == Some(&wire::StopReason::EndTurn)) + ); + assert!( + recording + .updates + .lock() + .unwrap() + .iter() + .all(|update| update.session_id == session_id) + ); + assert_eq!( + crate::session::branch::lookup_committed( + root.path(), + &child_id, + "failed-first-turn", + "edited-request", + ) + .unwrap() + .unwrap() + .session_id, + child_id + ); + + handle.close(); + let (reply, response) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + response.await.unwrap(); + timeout(Duration::from_secs(5), actor) + .await + .unwrap() + .unwrap(); + } + } + + enum GatedTurn { + Branch, + Prompt, + Autonomous, + } + + async fn cancelled_gated_branch_activation(close: bool, turn: GatedTurn) { + let initial_branch = matches!(turn, GatedTurn::Branch); + let autonomous = matches!(turn, GatedTurn::Autonomous); + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let source_id = crate::session::new_id(); + let child_id = crate::session::new_id(); + let prefix = vec![Item::text(ItemKind::System, "system")]; + let prompt = Item::text(ItemKind::User, "cancelled edited prompt"); + let selection = + crate::provider::ModelSelection::new(ProviderKind::OpenAiSubscription, "gpt-5.4"); + let checkout = "gated-cancel-checkout"; + let request = prompt_branches::submitted_request_id(&source_id, "cancelled edited prompt"); + let initial = crate::session::branch::prepare( + prefix.clone(), + source_id.clone(), + crate::session::branch::Boundary::new(0, &prefix).unwrap(), + checkout.into(), + crate::session::branch::SubmittedRequest { + id: request.clone(), + selection: crate::session::branch::CapturedSelection::new(&selection, None), + }, + prompt.clone(), + ) + .unwrap(); + let child = + crate::session::open_uncommitted(root.path(), &child_id, false, initial).unwrap(); + crate::session::branch::commit(&child.observer, &child.transcript).unwrap(); + drop(child); + let committed = crate::session::branch::load_history(root.path(), &child_id).unwrap(); + let session_id = wire::SessionId::new(child_id.clone()); + let integration = Arc::new(AcpIntegration::default()); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(child_id.clone()), + sink.clone(), + )) + .unwrap(); + handle.prepare_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let turns = Arc::new(AtomicU64::new(0)); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .task_manager(manager) + .transcript(prefix) + .input(if initial_branch { vec![prompt] } else { vec![] }) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(child_id.clone())).without_cache()) + .await + .unwrap(); + let driver = input_settlement.wrap(driver); + let mcp = crate::tools::mcp::empty(); + let mcp_events = mcp.subscribe(child_id.clone()); + let busy = Arc::new(AtomicBool::new(!autonomous)); + let (commands, receiver) = mpsc::channel(8); + let actor = SessionActor { + initial_generation: initial_branch.then_some(generation), + admission_released: Arc::new(Notify::new()), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(), + catalog: vec![], + commands: receiver, + mcp_events, + }; + // Use the real diagnostic I/O boundary to hold an admitted autonomous + // turn before its first step. No production callback or test field is + // needed: busy is the actual admission state, and the cancellation + // handle is the same interface used by session/cancel. + let cancellation = if autonomous { + assert!(crate::events::enabled()); + let (locked, ready) = oneshot::channel(); + let busy = busy.clone(); + let handle = handle.clone(); + let cancellation = std::thread::spawn(move || { + let stderr = std::io::stderr().lock(); + locked.send(()).unwrap(); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while !busy.load(Ordering::Acquire) { + assert!( + std::time::Instant::now() < deadline, + "autonomous admission never arrived" + ); + std::thread::yield_now(); + } + handle.interrupt(); + drop(stderr); + }); + ready.await.unwrap(); + Some(cancellation) + } else { + None + }; + let (activation, activated) = oneshot::channel(); + let task = tokio::spawn(async move { + activated.await.unwrap(); + session_actor(actor).await; + }); + // For ordinary prompts, open actor activation but hold the separate + // prompt response gate after input has been submitted to the driver. + let activation = if initial_branch || autonomous { + activation + } else { + activation.send(()).unwrap(); + let (reply, response) = oneshot::channel(); + commands + .send(Command::Prompt(PromptCommand { + request: wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "cancelled ordinary prompt", + ))], + ), + cancellation_generation: generation, + reply, + })) + .await + .unwrap(); + response.await.unwrap().unwrap() + }; + // The initial turn cannot run until its response gate opens. + let acknowledged = if close { + handle.close(); + let (reply, acknowledged) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + Some(acknowledged) + } else { + if !autonomous { + handle.interrupt(); + } + None + }; + // Queue a wake before activation: it must not resurrect the initial input. + mcp.publish( + &child_id, + crate::tools::mcp::McpEvent { + message: "late wake".into(), + }, + ); + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert_eq!(busy.load(Ordering::Acquire), !autonomous); + activation.send(()).unwrap(); + timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + if let Some(cancellation) = cancellation { + cancellation.join().unwrap(); + } + if let Some(acknowledged) = acknowledged { + acknowledged.await.unwrap(); + } + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert!( + handle.cancellation_handle().is_cancelled_since(generation), + "cancellation must advance the captured generation" + ); + assert!(!busy.load(Ordering::Acquire)); + assert!( + commands.is_closed(), + "cancelled pending input must lose its execution owner" + ); + assert!( + !recording + .updates + .lock() + .unwrap() + .iter() + .any(|update| matches!( + update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + )), + "an unstarted turn must not report model activity" + ); + // A lost response/retry discovers the same committed child, never a new + // destination or a second activation. Reload its full history passively. + let found = + crate::session::branch::find_committed(root.path(), &source_id, checkout, &request) + .unwrap() + .unwrap(); + assert_eq!(found.session_id, child_id); + assert_eq!( + crate::session::branch::load_history(root.path(), &child_id).unwrap(), + committed + ); + let input_settlement = crate::runtime::InputSettlement::default(); + let reloaded = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .transcript(found.transcript) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(child_id)).without_cache()) + .await + .unwrap(); + let mut reloaded = input_settlement.wrap(reloaded); + assert!(matches!( + reloaded.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + assert_eq!(turns.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn branch_cancel_before_activation_discards_pending_execution_but_preserves_retry() { + cancelled_gated_branch_activation(false, GatedTurn::Branch).await; + } + + #[tokio::test] + async fn branch_close_before_activation_discards_pending_execution_but_preserves_retry() { + cancelled_gated_branch_activation(true, GatedTurn::Branch).await; + } + + #[tokio::test] + async fn ordinary_prompt_cancel_before_response_gate_retires_pending_execution() { + cancelled_gated_branch_activation(false, GatedTurn::Prompt).await; + } + + #[tokio::test] + async fn autonomous_cancel_before_first_step_retires_pending_execution() { + // Isolate the process-global stderr lock and opt-in event setting from + // parallel tests. The child exercises the real actor and durable retry. + const CHILD: &str = "KIT_TEST_AUTONOMOUS_CANCEL_CHILD"; + if std::env::var_os(CHILD).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "protocols::acp::v2::tests::autonomous_cancel_before_first_step_retires_pending_execution", "--nocapture"]) + .env(CHILD, "1") + .env("KIT_RUNTIME_EVENTS", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("1 passed")); + return; + } + cancelled_gated_branch_activation(false, GatedTurn::Autonomous).await; + } + + #[test] + fn resume_location_rejects_invalid_routes_before_loaded_lookup() { + let root = tempfile::tempdir().unwrap(); + let other = tempfile::tempdir().unwrap(); + let canonical = crate::resilient_fs::canonicalize(root.path()).unwrap(); + assert!(validate_resume_location(&canonical, root.path(), false).is_ok()); + assert!(validate_resume_location(&canonical, root.path(), true).is_err()); + assert!(validate_resume_location(&canonical, other.path(), false).is_err()); + assert!(validate_resume_location(&canonical, &root.path().join("missing"), false).is_err()); + } + + #[tokio::test] + async fn every_active_turn_restores_source_route_before_activity_diagnostics() { + const CHILD: &str = "KIT_TEST_DIAGNOSTIC_ROUTE_CHILD"; + if std::env::var_os(CHILD).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "protocols::acp::v2::tests::every_active_turn_restores_source_route_before_activity_diagnostics", "--nocapture"]) + .env(CHILD, "1") + .env("KIT_RUNTIME_EVENTS", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8_lossy(&output.stderr); + let mut route = None; + let mut routes = Vec::new(); + let mut running = 0; + for line in stderr.lines() { + if let Some(crate::events::RuntimeEvent::SessionStarted { session_id }) = + crate::events::parse(line) + { + routes.push(session_id.clone()); + route = Some(session_id); + } else if let Ok(notification) = + serde_json::from_str::(line) + { + assert_eq!( + route.as_deref(), + Some(notification.session_id.to_string().as_str()) + ); + if matches!( + notification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ) { + running += 1; + } + } + } + assert_eq!(routes, ["failed-child", "source", "failed-child", "source"]); + assert_eq!(running, 2); + return; + } + // This test sink is a real serialized notification boundary. Sharing + // stderr with the runtime event stream makes routing order observable. + #[derive(Clone)] + struct DiagnosticSink(RecordingSink); + #[async_trait] + impl AcpSessionUpdateSink for DiagnosticSink { + fn update( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + eprintln!("{}", serde_json::to_string(¬ification).unwrap()); + self.0.update(notification) + } + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + self.update(notification) + } + async fn flush(&self) -> Result<(), AcpRuntimeError> { + self.0.flush().await + } + } + for origin in [ExecutionOrigin::Prompt, ExecutionOrigin::Autonomous] { + let integration = AcpIntegration::default(); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(DiagnosticSink(recording.clone())); + let session_id = wire::SessionId::new("source"); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new("source"), + sink.clone(), + )) + .unwrap(); + let activity = native_activity(session_id.clone(), sink.clone()); + let turns = Arc::new(AtomicU64::new(0)); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .observer(ResponseReplacementObserver::new( + integration.clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .input(vec![Item::text(ItemKind::User, "continue source")]) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new("source")).without_cache()) + .await + .unwrap(); + let mut driver = input_settlement.wrap(driver); + handle.prepare_injection_turn(); + let generation = handle.cancellation_handle().generation(); + handle.start_injection_turn(); + // Esc can return to source without an explicit resume. + establish_diagnostic_route(&wire::SessionId::new("failed-child")); + run_active_turn( + &session_id, + &integration, + &handle, + &mut driver, + &sink, + generation, + None, + &activity, + origin, + ) + .await + .unwrap(); + assert_eq!(turns.load(Ordering::Relaxed), 1); + assert!( + recording + .updates + .lock() + .unwrap() + .iter() + .all(|update| update.session_id == session_id) + ); + assert!( + recording + .updates + .lock() + .unwrap() + .iter() + .any(|update| matches!( + &update.update, + wire::SessionUpdate::AgentMessage(_) + | wire::SessionUpdate::AgentMessageChunk(_) + )) + ); + } + } + + #[test] + fn branch_admission_rejects_busy_and_pending_injections_without_leaking() { + let busy = Arc::new(AtomicBool::new(true)); + let work = Arc::new(Mutex::new(InjectionWork::default())); + assert!(BranchAdmission::claim(busy.clone(), work.clone()).is_err()); + assert!(!work.lock().unwrap().branch_reserved); + busy.store(false, Ordering::Release); + work.lock().unwrap().admitting = 1; + assert!(BranchAdmission::claim(busy.clone(), work.clone()).is_err()); + drop(InjectionAdmission(work.clone())); + let id = wire::MessageId::new("accepted-but-not-delivered"); + work.lock().unwrap().pending.insert(id.clone()); + assert!(BranchAdmission::claim(busy.clone(), work.clone()).is_err()); + assert!(!busy.load(Ordering::Acquire)); + work.lock().unwrap().pending.remove(&id); + let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); + assert!(busy.load(Ordering::Acquire)); + assert!(work.lock().unwrap().branch_reserved); + drop(admission); + assert!(!busy.load(Ordering::Acquire)); + assert!(!work.lock().unwrap().branch_reserved); + } + + #[test] + fn branch_injection_tracking_retains_only_accepted_work() { + let work = Arc::new(Mutex::new(InjectionWork::default())); + let failed = wire::MessageId::new("failed-receipt"); + let retained = wire::MessageId::new("accepted"); + work.lock() + .unwrap() + .pending + .extend([failed.clone(), retained.clone()]); + drop(TrackedInjection { + work: work.clone(), + id: failed.clone(), + retained: false, + }); + drop(TrackedInjection { + work: work.clone(), + id: retained.clone(), + retained: true, + }); + assert_eq!(work.lock().unwrap().pending, HashSet::from([retained])); + } + + #[tokio::test] + async fn branch_actor_reservation_blocks_following_work_until_release() { + let busy = Arc::new(AtomicBool::new(false)); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); + let (reply, response) = oneshot::channel(); + let next_command = Arc::new(AtomicBool::new(false)); + let processed = next_command.clone(); + let actor = tokio::spawn(async move { + hold_branch_reservation(42, admission, reply).await; + processed.store(true, Ordering::Release); + }); + let (checkout, release) = response.await.unwrap().unwrap(); + assert_eq!(checkout, 42); + tokio::task::yield_now().await; + assert!(!next_command.load(Ordering::Acquire)); + assert!(busy.load(Ordering::Acquire)); + assert!(work.lock().unwrap().branch_reserved); + drop(release); // failed submission/cancelled receiver also unblocks + actor.await.unwrap(); + assert!(next_command.load(Ordering::Acquire)); + assert!(!busy.load(Ordering::Acquire)); + assert!(!work.lock().unwrap().branch_reserved); + } + + #[tokio::test] + async fn branch_actor_failed_reply_releases_admission() { + let busy = Arc::new(AtomicBool::new(false)); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); + let (reply, response) = oneshot::channel(); + drop(response); + hold_branch_reservation((), admission, reply).await; + assert!(!busy.load(Ordering::Acquire)); + assert!(!work.lock().unwrap().branch_reserved); + } + + #[derive(Clone, Copy)] + enum RacingBranch { + Prepare, + Reserve, + Abandon, + } + + async fn selected_completion_survives_branch_admission(background: bool, branch: RacingBranch) { + use agentkit_core::{TaskId, ToolCallId, TurnId}; + use agentkit_task_manager::{ContinuePolicy, TaskLaunchRequest, TaskStartContext}; + use agentkit_tools_core::{ + AllowAllPermissions, BasicToolExecutor, OwnedToolContext, ToolRequest, ToolSource, + }; + + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let session_id = wire::SessionId::new("selected-completion"); + let loop_id = SessionId::new("selected-completion"); + let integration = Arc::new(AcpIntegration::default()); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let generations = Arc::new(Mutex::new(Vec::new())); + let activity = SessionActivity::new({ + let generations = generations.clone(); + move |transition| { + generations + .lock() + .unwrap() + .push((transition.id, transition.active)); + Ok(()) + } + }); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + let cancellation_generation = handle.cancellation_handle().generation(); + let task_manager = + AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background); + let tasks = task_manager.handle(); + let release_task = Arc::new(Notify::new()); + if background { + let tools = ToolRegistry::new().with(BlockingTool { + spec: ToolSpec { + name: ToolName::new("race-completion"), + description: "controlled background completion".into(), + input_schema: json!({"type": "object"}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered: Arc::new(AtomicBool::new(false)), + release: release_task.clone(), + }); + let task_id = TaskId::new("race-task"); + task_manager + .start_task( + TaskLaunchRequest::plain( + Some(task_id.clone()), + ToolRequest { + call_id: ToolCallId::new("race-call"), + tool_name: ToolName::new("race-completion"), + input: json!({}), + session_id: loop_id.clone(), + turn_id: TurnId::new("background-turn"), + metadata: MetadataMap::new(), + }, + ), + TaskStartContext { + executor: Arc::new(BasicToolExecutor::new([ + Arc::new(tools) as Arc + ])), + tool_context: OwnedToolContext { + session_id: loop_id.clone(), + turn_id: TurnId::new("background-turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }, + }, + ) + .await + .unwrap(); + tasks + .set_continue_policy(task_id, ContinuePolicy::RequestContinue) + .await + .unwrap(); + } + let turns = Arc::new(AtomicU64::new(0)); + let observer = ResponseReplacementObserver::new( + integration.as_ref().clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + ); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .observer(observer) + .task_manager(task_manager) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let driver = input_settlement.wrap(driver); + let mcp = crate::tools::mcp::empty(); + let mcp_events = mcp.subscribe(session_id.to_string()); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let busy = Arc::new(AtomicBool::new(false)); + let (commands, receiver) = mpsc::channel(8); + let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); + let completion = tasks.clone(); + let mut actor = Box::pin(session_actor(SessionActor { + initial_generation: None, + admission_released: work.lock().unwrap().admission_released.clone(), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(), + catalog: vec![], + commands: receiver, + mcp_events, + })); + if background { + release_task.notify_one(); + timeout(Duration::from_secs(5), completion.wait_for_idle()) + .await + .unwrap(); + } else { + mcp.publish( + &session_id.to_string(), + crate::tools::mcp::McpEvent { + message: "selected MCP completion".into(), + }, + ); + } + assert!(futures_util::poll!(&mut actor).is_pending()); + // The actor consumed all ready events and is now suspended behind the + // real admission owner. Starting it as a task cannot admit that wake. + let actor = tokio::spawn(actor); + assert!(busy.load(Ordering::Acquire)); + assert_eq!(turns.load(Ordering::Relaxed), 0); + match branch { + RacingBranch::Prepare => { + let (reply, response) = oneshot::channel(); + commands + .send(Command::PreparePromptBranch { + address: "unused: unsettled must reject first".into(), + admission, + reply, + }) + .await + .unwrap(); + let error = timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .expect_err("selected completion must reject checkout preparation"); + assert!(error.to_string().contains("unsettled work"), "{error}"); + } + RacingBranch::Reserve => { + let (reply, response) = oneshot::channel(); + commands + .send(Command::ReservePromptBranch { + checkout_token: "unused: unsettled must reject first".into(), + admission, + reply, + }) + .await + .unwrap(); + let error = timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .err() + .expect("selected completion must reject child reservation"); + assert!(error.to_string().contains("unsettled work"), "{error}"); + } + RacingBranch::Abandon => { + // A read-only barrier proves the event handler returned while + // busy is still held. No branch command will wake this actor. + let (reply, response) = oneshot::channel(); + commands.send(Command::Snapshot { reply }).await.unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(turns.load(Ordering::Relaxed), 0); + drop(admission); + } + } + // No prompt, snapshot, MCP event or task event is sent to wake the drive. + timeout(Duration::from_secs(5), async { + loop { + if generations + .lock() + .unwrap() + .last() + .is_some_and(|(_, active)| !active) + && !busy.load(Ordering::Acquire) + { + break; } - }, - agent_client_protocol::on_receive_request!(), + tokio::task::yield_now().await; + } + }) + .await + .expect("selected completion was stranded after branch admission released"); + let (reply, response) = oneshot::channel(); + commands.send(Command::Snapshot { reply }).await.unwrap(); + let snapshot = response.await.unwrap().unwrap(); + let transcript = serde_json::to_string(&snapshot.canonical_transcript).unwrap(); + assert_eq!( + transcript + .matches(if background { + "background done" + } else { + "selected MCP completion" + }) + .count(), + 1 ); - Ok(agent) -} - -#[cfg(test)] -mod tests { - use std::{ - collections::VecDeque, - sync::{atomic::AtomicUsize, mpsc as std_mpsc}, - }; + assert_eq!(transcript.matches("autonomous content").count(), 1); + assert_eq!(turns.load(Ordering::Relaxed), 1); + assert_eq!(*generations.lock().unwrap(), vec![(1, true), (1, false)]); + assert_eq!( + handle.cancellation_handle().generation(), + cancellation_generation + ); + assert_eq!(recording.flushes.load(Ordering::Relaxed), 1); + assert!(!busy.load(Ordering::Acquire)); + assert!(!work.lock().unwrap().branch_reserved); + let (reply, response) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + response.await.unwrap(); + actor.await.unwrap(); + assert_eq!(turns.load(Ordering::Relaxed), 1); + } - use serde_json::json; + #[tokio::test] + async fn branch_actor_selected_mcp_completion_is_not_stranded() { + for branch in [ + RacingBranch::Prepare, + RacingBranch::Reserve, + RacingBranch::Abandon, + ] { + selected_completion_survives_branch_admission(false, branch).await; + } + } - use agent_client_protocol::schema::MaybeUndefined; - use agentkit_core::{DataRef, MetadataMap, Modality, TurnCancellation}; - use agentkit_loop::{ - Agent, ModelAdapter, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, - TurnRequest, TurnResult, - }; - use agentkit_task_manager::{AsyncTaskManager, RoutingDecision, TaskManager}; - use agentkit_tools_core::{ToolAnnotations, ToolName, ToolRegistry, ToolSpec}; - use tokio::{ - sync::Notify, - time::{Duration, timeout}, - }; + #[tokio::test] + async fn branch_actor_selected_background_completion_is_not_stranded() { + for branch in [ + RacingBranch::Prepare, + RacingBranch::Reserve, + RacingBranch::Abandon, + ] { + selected_completion_survives_branch_admission(true, branch).await; + } + } - use super::*; - use crate::protocols::acp::tests::{BlockingTool, ScriptAdapter}; + #[tokio::test] + async fn branch_settled_snapshot_rejects_input_mcp_and_unacknowledged_work() { + let (mut driver, turns) = test_driver(TestOutcome::Content, "branch-settled").await; + let tasks = AsyncTaskManager::new().handle(); + let jobs = BackgroundJobs::default(); + let mcp = crate::tools::mcp::empty(); + let mut events = mcp.subscribe("branch-settled".into()); + let original = driver.snapshot().transcript; + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .unwrap(); + mcp.publish( + "branch-settled", + crate::tools::mcp::McpEvent { + message: "queued update".into(), + }, + ); + assert!( + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .is_err() + ); + assert!(events.has_pending()); // checking never consumes the event + events.recv().await.unwrap(); + jobs.register_foreground_for_test("branch-job"); + assert!( + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .is_err() + ); + jobs.detach("branch-job"); + jobs.finish_for_test("branch-job"); + assert!(jobs.activity().unacknowledged_terminals); + assert!( + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .is_err() + ); + jobs.acknowledge_terminal(&agentkit_core::ToolCallId::new("branch-job")); + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "queued prompt")]) + .unwrap(); + assert!( + settled_branch_source(&driver, &tasks, &jobs, &events) + .await + .is_err() + ); + assert_eq!(driver.snapshot().transcript, original); + assert_eq!(driver.snapshot().pending_input.len(), 1); + assert_eq!(turns.load(Ordering::Relaxed), 0); + } fn terminal_auth_initialize_request() -> wire::InitializeRequest { wire::InitializeRequest::new( @@ -2732,6 +4972,8 @@ mod tests { Content, FinishError, ProviderError, + FinishErrorThenContent, + ProviderErrorThenContent, } struct TestAdapter { @@ -2842,7 +5084,14 @@ mod tests { if let Some(handle) = self.interrupt.take() { handle.interrupt(); } - match self.outcome { + let outcome = self.outcome; + if matches!( + outcome, + TestOutcome::FinishErrorThenContent | TestOutcome::ProviderErrorThenContent + ) { + self.outcome = TestOutcome::Content; + } + match outcome { TestOutcome::ToolThenContent => { self.outcome = TestOutcome::Content; let call = agentkit_core::ToolCallPart::new( @@ -2890,7 +5139,7 @@ mod tests { ]), }) } - TestOutcome::FinishError => Ok(TestTurn { + TestOutcome::FinishError | TestOutcome::FinishErrorThenContent => Ok(TestTurn { events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult { model: None, response_id: None, @@ -2900,7 +5149,9 @@ mod tests { metadata: MetadataMap::new(), })]), }), - TestOutcome::ProviderError => Err(LoopError::Provider("provider failed".into())), + TestOutcome::ProviderError | TestOutcome::ProviderErrorThenContent => { + Err(LoopError::Provider("provider failed".into())) + } } } } @@ -2928,7 +5179,9 @@ mod tests { interrupt: Option, ) -> (LoopDriver, Arc) { let turns = Arc::new(AtomicU64::new(0)); + let input_settlement = crate::runtime::InputSettlement::default(); let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome, turns: Arc::clone(&turns), @@ -2939,6 +5192,7 @@ mod tests { .start(SessionConfig::new(SessionId::new(session_id)).without_cache()) .await .unwrap(); + let driver = input_settlement.wrap(driver); (driver, turns) } @@ -2959,7 +5213,9 @@ mod tests { .unwrap(); let observer = ResponseReplacementObserver::new(integration, sink, session_id, activity.clone()); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(StreamingCancellationAdapter { interrupt: handle.clone(), }) @@ -2970,6 +5226,7 @@ mod tests { .start(SessionConfig::new(loop_session_id).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::text(ItemKind::User, "cancel")]) .unwrap(); @@ -3131,7 +5388,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::ProviderError, turns: turns.clone(), @@ -3148,6 +5407,7 @@ mod tests { ) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); let (reply, response) = oneshot::channel(); let command = PromptCommand { request: wire::PromptRequest::new( @@ -3256,7 +5516,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(ScriptAdapter { turns: Arc::clone(&turns), user_items_seen: Arc::new(AtomicUsize::new(0)), @@ -3270,6 +5532,7 @@ mod tests { .start(SessionConfig::new(SessionId::new("v2-structured-loop")).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::text(ItemKind::User, "start background")]) .unwrap(); @@ -3356,7 +5619,7 @@ mod tests { fn bind_test_session( integration: &AcpIntegration, session_id: &wire::SessionId, - sink: RecordingSink, + sink: impl AcpSessionUpdateSink + 'static, ) -> AcpSessionHandle { let handle = integration .bind_session(AcpSessionBinding::new( @@ -3457,7 +5720,8 @@ mod tests { async fn cancellation_before_tool_boundary_retires_turn() { let integration = AcpIntegration::default(); let session_id = wire::SessionId::new("boundary-cancel"); - let handle = bind_test_session(&integration, &session_id, RecordingSink::default()); + let sink = ResponseReplacementSink::new(RecordingSink::default()); + let handle = bind_test_session(&integration, &session_id, sink.clone()); let generation = handle.cancellation_handle().generation(); let (mut driver, _) = test_driver_with_interrupt( TestOutcome::ToolThenContent, @@ -3469,9 +5733,16 @@ mod tests { .submit_input(vec![Item::text(ItemKind::User, "first")]) .unwrap(); assert_eq!( - drive_prompt(&session_id, &mut driver, &handle, generation, None) - .await - .unwrap(), + drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements + ) + .await + .unwrap(), FinishReason::Cancelled ); assert!( @@ -3488,9 +5759,16 @@ mod tests { .unwrap(); let generation = handle.cancellation_handle().generation(); assert_eq!( - drive_prompt(&session_id, &mut driver, &handle, generation, None) - .await - .unwrap(), + drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements + ) + .await + .unwrap(), FinishReason::Completed ); assert!( @@ -3509,18 +5787,20 @@ mod tests { let integration = AcpIntegration::default(); let session_id = wire::SessionId::new("injection-cancel"); let recording = RecordingSink::default(); - let handle = bind_test_session(&integration, &session_id, recording.clone()); + let sink = ResponseReplacementSink::new(recording.clone()); + let handle = bind_test_session(&integration, &session_id, sink.clone()); let (receipt, _client, server) = staged_injection(integration.clone(), &session_id).await; let message_id = receipt.message_id().clone(); - let sink = ResponseReplacementSink::new(recording.clone()); let activity = native_activity(session_id.clone(), sink.clone()); let observer = ResponseReplacementObserver::new( integration, - sink, + sink.clone(), session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome, turns: Arc::new(AtomicU64::new(0)), @@ -3532,6 +5812,7 @@ mod tests { .start(SessionConfig::new(SessionId::new("injection-cancel")).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::text(ItemKind::User, "first")]) .unwrap(); @@ -3539,7 +5820,14 @@ mod tests { { let turn = activity.execute( ExecutionOrigin::Prompt, - drive_prompt(&session_id, &mut driver, &handle, generation, None), + drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements, + ), |reason| Some(reason.clone()), ); tokio::pin!(turn); @@ -3579,7 +5867,14 @@ mod tests { activity .execute( ExecutionOrigin::Prompt, - drive_prompt(&session_id, &mut driver, &handle, generation, None), + drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements + ), |reason| Some(reason.clone()) ) .await @@ -3613,7 +5908,8 @@ mod tests { let integration = AcpIntegration::default(); let session_id = wire::SessionId::new("terminal-inject"); let recording = RecordingSink::default(); - let handle = bind_test_session(&integration, &session_id, recording.clone()); + let sink = ResponseReplacementSink::new(recording.clone()); + let handle = bind_test_session(&integration, &session_id, sink.clone()); let (receipt, _client, server) = staged_injection(integration, &session_id).await; let message_id = receipt.message_id().clone(); let (mut driver, _) = test_driver(TestOutcome::Content, "terminal-inject").await; @@ -3622,7 +5918,14 @@ mod tests { .unwrap(); let generation = handle.cancellation_handle().generation(); { - let turn = drive_prompt(&session_id, &mut driver, &handle, generation, None); + let turn = drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements, + ); tokio::pin!(turn); assert!(futures_util::poll!(&mut turn).is_pending()); receipt.activate_after_response().await.unwrap(); @@ -3669,14 +5972,23 @@ mod tests { let integration = AcpIntegration::default(); let session_id = wire::SessionId::new("finish-error"); let recording = RecordingSink::default(); - let handle = bind_test_session(&integration, &session_id, recording.clone()); + let sink = ResponseReplacementSink::new(recording.clone()); + let handle = bind_test_session(&integration, &session_id, sink.clone()); let (receipt, _client, server) = staged_injection(integration, &session_id).await; let (mut driver, _) = test_driver(TestOutcome::FinishError, "finish-error").await; driver .submit_input(vec![Item::text(ItemKind::User, "fail")]) .unwrap(); let generation = handle.cancellation_handle().generation(); - let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; + let result = drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements, + ) + .await; assert!( matches!(result, Err(AcpRuntimeError::Loop(message)) if message == "model turn failed") ); @@ -3694,13 +6006,13 @@ mod tests { #[tokio::test] async fn cancellation_race_wins_over_provider_error() { let integration = AcpIntegration::default(); - let sink = RecordingSink::default(); + let sink = ResponseReplacementSink::new(RecordingSink::default()); let session_id = wire::SessionId::new("cancel-race"); let handle = integration .bind_session(AcpSessionBinding::new( session_id.clone(), SessionId::new("cancel-race-loop"), - sink, + sink.clone(), )) .unwrap(); handle.prepare_injection_turn(); @@ -3713,7 +6025,15 @@ mod tests { .submit_input(vec![Item::text(ItemKind::User, "cancel")]) .unwrap(); - let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; + let result = drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements, + ) + .await; assert_eq!(result.unwrap(), FinishReason::Cancelled); } @@ -3722,13 +6042,22 @@ mod tests { async fn provider_error_without_cancellation_remains_an_error() { let integration = AcpIntegration::default(); let session_id = wire::SessionId::new("provider-error"); - let handle = bind_test_session(&integration, &session_id, RecordingSink::default()); + let sink = ResponseReplacementSink::new(RecordingSink::default()); + let handle = bind_test_session(&integration, &session_id, sink.clone()); let (mut driver, _) = test_driver(TestOutcome::ProviderError, "provider-error").await; driver .submit_input(vec![Item::text(ItemKind::User, "fail")]) .unwrap(); let generation = handle.cancellation_handle().generation(); - let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; + let result = drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + None, + &sink.acknowledgements, + ) + .await; assert!(matches!(result, Err(AcpRuntimeError::Loop(_)))); } @@ -3753,7 +6082,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::Content, turns: turns.clone(), @@ -3770,6 +6101,7 @@ mod tests { ) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); let busy = AtomicBool::new(false); drive_autonomous( @@ -3812,7 +6144,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::Content, turns: Arc::clone(&turns), @@ -3824,6 +6158,7 @@ mod tests { .start(SessionConfig::new(loop_session_id).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::notification("background event")]) .unwrap(); @@ -3890,7 +6225,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::ProviderError, turns: turns.clone(), @@ -3907,6 +6244,7 @@ mod tests { ) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::notification("background event")]) .unwrap(); @@ -3964,7 +6302,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::ProviderError, turns: turns.clone(), @@ -3981,6 +6321,7 @@ mod tests { ) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::notification("background event")]) .unwrap(); @@ -4028,7 +6369,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::Content, turns: Arc::new(AtomicU64::new(0)), @@ -4040,6 +6383,7 @@ mod tests { .start(SessionConfig::new(SessionId::new("failed-flush-loop")).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::notification("work")]) .unwrap(); @@ -4090,7 +6434,9 @@ mod tests { session_id.clone(), activity.clone(), ); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::FinishError, turns: turns.clone(), @@ -4107,6 +6453,7 @@ mod tests { ) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); driver .submit_input(vec![Item::notification("background event")]) .unwrap(); @@ -4191,6 +6538,7 @@ mod tests { server.sessions.lock().unwrap().insert( session_id.clone(), SessionHandle { + injections: Arc::new(Mutex::new(InjectionWork::default())), token: 1, commands, integration, @@ -4332,6 +6680,7 @@ mod tests { completed: completion, session_id, session: SessionHandle { + injections: Arc::new(Mutex::new(InjectionWork::default())), token, commands, integration, @@ -4651,6 +7000,7 @@ mod tests { completed: completion, session_id: published_session_id, session: SessionHandle { + injections: Arc::new(Mutex::new(InjectionWork::default())), token, commands, integration, @@ -4836,6 +7186,212 @@ mod tests { result.expect("exhaustion client timed out").unwrap(); } + #[tokio::test] + async fn branch_router_replays_loaded_source_and_durable_child_without_execution() { + let root = tempfile::tempdir().unwrap(); + let workspace = root.path().to_path_buf(); + let source_id = crate::session::new_id(); + let source = crate::session::open( + root.path(), + &source_id, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "original prompt"), + Item::text(ItemKind::Assistant, "original answer"), + ], + ) + .unwrap(); + let prefix = source.transcript[..1].to_vec(); + drop(source); + let source_before = crate::session::branch::load_history(root.path(), &source_id).unwrap(); + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + ) + .unwrap(); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime, SessionRegistry::new()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let updates = Arc::new(Mutex::new(Vec::::new())); + let received = updates.clone(); + let result = agent_client_protocol::Client + .v2() + .on_receive_notification( + async move |update: wire::UpdateSessionNotification, _cx| { + received.lock().unwrap().push(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("checkout-test", "0"), + )) + .block_task() + .await?; + cx.send_request( + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ) + .block_task() + .await?; + let listed = cx + .send_request(ListPromptBranchesRequest { + session_id: wire::SessionId::new(source_id.clone()), + }) + .block_task() + .await?; + assert_eq!(listed.boundaries.len(), 1); + let address = listed.boundaries[0].address.clone(); + // An integration-rejected injection must release the admission + // tracker, so a subsequent settled checkout still succeeds. + cx.send_request(wire::InjectSessionRequest::new( + source_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "queued steer", + ))], + )) + .block_task() + .await + .expect_err("idle session cannot accept a new injection"); + let prepared = cx + .send_request(PreparePromptBranchRequest { + session_id: wire::SessionId::new(source_id.clone()), + address, + }) + .block_task() + .await?; + assert_eq!(prepared.original_text, "original prompt"); + // Simulate the durable commit from a prior successful submission. + // Recovery must never dispatch a provider turn, even if the live + // source revision has subsequently changed. + let child_id = crate::session::new_id(); + let selection = crate::provider::ModelSelection::new( + crate::ProviderKind::OpenRouter, + "test-model", + ); + let initial = crate::session::branch::prepare( + prefix.clone(), + source_id.clone(), + crate::session::branch::Boundary::new(0, &prefix).unwrap(), + prepared.checkout_token.clone(), + crate::session::branch::SubmittedRequest { + id: prompt_branches::submitted_request_id(&source_id, "edited prompt"), + selection: crate::session::branch::CapturedSelection::new(&selection, None), + }, + Item::text(ItemKind::User, "edited prompt"), + ) + .unwrap(); + let child = crate::session::open_uncommitted(&workspace, &child_id, false, initial) + .unwrap(); + crate::session::branch::commit(&child.observer, &child.transcript).unwrap(); + let mut child_future = child.transcript.clone(); + child_future.push(Item::text(ItemKind::Assistant, "existing child answer")); + child.observer.replace(&child_future).unwrap(); + drop(child); + let child_before = + crate::session::branch::load_history(&workspace, &child_id).unwrap(); + cx.send_request(wire::SetSessionConfigOptionRequest::new( + source_id.clone(), + "reasoning_effort", + "high", + )) + .block_task() + .await?; + let request = SubmitPromptBranchRequest { + session_id: wire::SessionId::new(source_id.clone()), + checkout_token: prepared.checkout_token, + text: "edited prompt".into(), + }; + let (first, second) = tokio::join!( + async { cx.send_request(request.clone()).block_task().await }, + async { cx.send_request(request.clone()).block_task().await }, + ); + assert_eq!(first?.session_id.to_string(), child_id); + assert_eq!(second?.session_id.to_string(), child_id); + let changed = cx + .send_request(wire::SetSessionConfigOptionRequest::new( + child_id.clone(), + "reasoning_effort", + "high", + )) + .block_task() + .await?; + let retry = cx.send_request(request.clone()).block_task().await?; + assert_eq!( + serde_json::to_value(retry.config_options).unwrap(), + serde_json::to_value(changed.config_options).unwrap() + ); + let mut different = request.clone(); + different.text = "different edit".into(); + cx.send_request(different) + .block_task() + .await + .expect_err("token must bind the original request"); + let resumed = cx + .send_request( + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ) + .block_task() + .await?; + assert!( + serde_json::to_string(&resumed.config_options) + .unwrap() + .contains("high") + ); + cx.send_request(wire::CloseSessionRequest::new(source_id.clone())) + .block_task() + .await?; + // Durable lookup remains available with no source actor at all. + assert_eq!( + cx.send_request(request) + .block_task() + .await? + .session_id + .to_string(), + child_id + ); + cx.send_request(wire::CloseSessionRequest::new(child_id.clone())) + .block_task() + .await?; + assert_eq!( + crate::session::branch::load_history(&workspace, &source_id).unwrap(), + source_before + ); + assert_eq!( + crate::session::branch::load_history(&workspace, &child_id).unwrap(), + child_before + ); + let received = updates.lock().unwrap(); + assert!(received.iter().any(|update| { + update.session_id.to_string() == child_id + && serde_json::to_string(&update.update) + .unwrap() + .contains("existing child answer") + })); + assert!(!received.iter().any(|update| matches!( + update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ))); + Ok(()) + }); + let result = timeout(Duration::from_secs(10), result).await; + server.abort(); + let _ = server.await; + result + .expect("checkout routing timed out") + .expect("checkout routing failed"); + } + #[tokio::test] async fn v2_router_advertises_and_routes_pending_injection_replacement() { let root = tempfile::tempdir().unwrap(); @@ -5195,6 +7751,341 @@ mod tests { } } + #[tokio::test] + async fn model_switch_cancel_before_first_step_retires_queued_prompt_and_mcp_wake() { + use std::io::{Read as _, Write as _}; + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + const CHILD: &str = "KIT_TEST_MODEL_SWITCH_CANCEL_CHILD"; + const ROUTE: &str = "compact-cancel-io-"; + if std::env::var_os(CHILD).is_none() { + // Like the autonomous regression, isolate process-global diagnostic + // I/O. Pipe backpressure gives a stronger handshake than observing + // busy: the route is emitted AFTER submit_input returns, and the + // actor cannot reach its first next() until this write completes. + let mut child = tokio::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "protocols::acp::v2::tests::model_switch_cancel_before_first_step_retires_queued_prompt_and_mcp_wake", "--nocapture"]) + .env(CHILD, "1") + .env("KIT_RUNTIME_EVENTS", "1") + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + let mut stderr = child.stderr.take().unwrap(); + let mut stdout = child.stdout.take().unwrap(); + let mut stdin = child.stdin.take().unwrap(); + timeout(Duration::from_secs(15), async { + let needle = format!("\"session_id\":\"{ROUTE}"); + let mut prefix = Vec::new(); + let mut chunk = [0; 128]; + loop { + let read = stderr.read(&mut chunk).await.unwrap(); + assert_ne!(read, 0, "child never reached the post-submission route"); + prefix.extend_from_slice(&chunk[..read]); + if prefix + .windows(needle.len()) + .any(|part| part == needle.as_bytes()) + { + break; + } + } + // Stop draining the oversized route until the child confirms + // the real interrupt AND both queued wake sources. This holds + // the actor at a genuine I/O boundary, not a test-only hook. + stdin.write_all(b"!").await.unwrap(); + let mut output = Vec::new(); + loop { + let byte = stdout.read_u8().await.unwrap(); + output.push(byte); + if byte == 1 { + break; + } + } + let mut diagnostics = Vec::new(); + stderr.read_to_end(&mut diagnostics).await.unwrap(); + stdout.read_to_end(&mut output).await.unwrap(); + let status = child.wait().await.unwrap(); + assert!( + status.success(), + "{}", + String::from_utf8_lossy(&diagnostics[diagnostics.len().saturating_sub(2000)..]) + ); + assert!(String::from_utf8_lossy(&output).contains("1 passed")); + }) + .await + .expect("model-switch cancellation child did not finish"); + return; + } + + assert!(crate::events::enabled()); + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let durable_id = crate::session::new_id(); + let opened = crate::session::open( + root.path(), + &durable_id, + false, + false, + vec![ + Item::text(ItemKind::User, "older content ".repeat(20_000)), + Item::text(ItemKind::Assistant, "recent response") + .with_usage(Usage::new(agentkit_core::TokenUsage::new(100, 0))), + ], + ) + .unwrap(); + let original = opened.transcript.clone(); + let path = crate::session::transcript_path_for_test(root.path(), &durable_id); + let bytes = std::fs::read(&path).unwrap(); + // The transport ID need not be the disk ID. Exceed the diagnostic pipe + // capacity so its first bytes are observable while emit still blocks. + let session_id = wire::SessionId::new(format!("{ROUTE}{}", "x".repeat(2 * 1024 * 1024))); + let loop_id = SessionId::new(durable_id.clone()); + let integration = Arc::new(AcpIntegration::default()); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + let selection = SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + // Only the model-boundary fake appends to seen; readers hold no lock + // across an await, I/O, cancellation, or actor cleanup. The production + // session observer owns all durable transcript writes/replacements. + let seen = Arc::new(Mutex::new(Vec::new())); + let turns = Arc::new(AtomicU64::new(0)); + let compactor = crate::compaction::automatic( + SwitchSummaryAdapter { + selection: selection.clone(), + seen: seen.clone(), + outcome: TestOutcome::Content, + interrupt: None, + }, + Default::default(), + Some(opened.observer.clone()), + loop_id.clone(), + ) + .unwrap(); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .task_manager(manager) + .mutator(compactor) + .observer(ResponseReplacementObserver::new( + (*integration).clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .transcript_observer(opened.observer) + .transcript(opened.transcript) + .cancellation(handle.cancellation_handle()) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let driver = input_settlement.wrap(driver); + let busy = Arc::new(AtomicBool::new(false)); + let (commands, receiver) = mpsc::channel(8); + let mcp = crate::tools::mcp::empty(); + let mcp_events = mcp.subscribe(durable_id.clone()); + let actor = tokio::spawn(session_actor(SessionActor { + initial_generation: None, + admission_released: Arc::new(Notify::new()), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: selection.clone(), + catalog: vec![crate::provider::ModelGroup { + provider: ProviderKind::OpenAiSubscription, + models: vec!["gpt-5.4-mini".into()], + context_windows: [("gpt-5.4-mini".into(), 150)].into_iter().collect(), + }], + commands: receiver, + mcp_events, + })); + let generation = handle.cancellation_handle().generation(); + let request = wire::SetSessionConfigOptionRequest::new( + session_id.clone(), + super::super::MODEL_CONFIG_ID, + "openai-subscription:gpt-5.4-mini", + ); + let (reply, response) = oneshot::channel(); + commands + .send(Command::SetConfig { + request: request.clone(), + reply, + cancellation_generation: generation, + }) + .await + .unwrap(); + let warning = response.await.unwrap().unwrap_err(); + let warning: model_switch::Warning = + serde_json::from_value(warning.data.unwrap()[model_switch::META].clone()).unwrap(); + assert_eq!(warning.guarded_tokens, "120"); + assert_eq!(warning.target_window, 150); + assert!(!busy.load(Ordering::Acquire)); + assert_eq!(selection.selection().unwrap().model, "gpt-5.4"); + assert!(seen.lock().unwrap().is_empty()); + assert_eq!(turns.load(Ordering::Relaxed), 0); + // A warning is not retirement: this confirmation must reach the SAME + // actor-owned Guard with its original transcript and generation. + let mut request = request; + request.meta = Some(serde_json::Map::from_iter([( + model_switch::META.into(), + serde_json::to_value(model_switch::Confirmation { + token: warning.token, + action: model_switch::Decision::Compact, + }) + .unwrap(), + )])); + let (queued_reply, queued_response) = oneshot::channel(); + let cancellation = { + let handle = handle.clone(); + let busy = busy.clone(); + let commands = commands.clone(); + let turns = turns.clone(); + let seen = seen.clone(); + let durable_id = durable_id.clone(); + std::thread::spawn(move || { + let mut signal = [0]; + std::io::stdin().read_exact(&mut signal).unwrap(); + assert_eq!(signal, *b"!"); + assert!(busy.load(Ordering::Acquire)); + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert!(seen.lock().unwrap().is_empty()); + handle.interrupt(); + assert!(handle.cancellation_handle().is_cancelled_since(generation)); + // Use a fresh generation: rejection cannot be explained by + // ordinary stale-prompt cancellation. Both wake sources are + // already queued when the diagnostic write is released. + commands + .try_send(Command::Prompt(PromptCommand { + request: wire::PromptRequest::new( + session_id, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "later prompt", + ))], + ), + cancellation_generation: handle.cancellation_handle().generation(), + reply: queued_reply, + })) + .unwrap(); + mcp.publish( + &durable_id, + crate::tools::mcp::McpEvent { + message: "later MCP wake".into(), + }, + ); + std::io::stdout().write_all(&[1]).unwrap(); + std::io::stdout().flush().unwrap(); + }) + }; + let (reply, response) = oneshot::channel(); + commands + .send(Command::SetConfig { + request, + reply, + cancellation_generation: generation, + }) + .await + .unwrap(); + let error = response.await.unwrap().unwrap_err(); + assert!( + error + .to_string() + .contains("compaction did not complete; model unchanged") + ); + assert!( + queued_response.await.is_err(), + "retirement must discard the queued prompt" + ); + timeout(Duration::from_secs(5), actor) + .await + .unwrap() + .unwrap(); + cancellation.join().unwrap(); + assert!(commands.is_closed()); + assert!(!busy.load(Ordering::Acquire)); + assert_eq!(selection.selection().unwrap().model, "gpt-5.4"); + assert!( + seen.lock().unwrap().is_empty(), + "cancelled compaction must never start its model" + ); + assert_eq!(turns.load(Ordering::Relaxed), 0); + assert!( + !recording + .updates + .lock() + .unwrap() + .iter() + .any(|update| matches!( + update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + )) + ); + assert_eq!(std::fs::read(path).unwrap(), bytes); + let transcript = crate::session::load(root.path(), &durable_id).unwrap(); + assert_eq!(transcript, original); + let input_settlement = crate::runtime::InputSettlement::default(); + let reloaded = Agent::builder() + .mutator(input_settlement.clone()) + .model(TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }) + .transcript(transcript) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(durable_id)).without_cache()) + .await + .unwrap(); + let mut reloaded = input_settlement.wrap(reloaded); + assert!(matches!( + reloaded.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + assert_eq!(turns.load(Ordering::Relaxed), 0); + } + #[tokio::test] async fn model_switch_compacts_with_original_model_before_selecting_and_keeps_it_on_failure_or_cancel() { @@ -5230,7 +8121,9 @@ mod tests { loop_id.clone(), ) .unwrap(); - let mut driver = Agent::builder() + let input_settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(TestAdapter { outcome: TestOutcome::Content, turns: Arc::new(AtomicU64::new(0)), @@ -5247,6 +8140,7 @@ mod tests { .start(SessionConfig::new(loop_id).without_cache()) .await .unwrap(); + let mut driver = input_settlement.wrap(driver); handle.prepare_injection_turn(); handle.start_injection_turn(); let generation = handle.cancellation_handle().generation(); diff --git a/src/protocols/acp/v2/tests/migration_resume.rs b/src/protocols/acp/v2/tests/migration_resume.rs new file mode 100644 index 0000000..08dfd7e --- /dev/null +++ b/src/protocols/acp/v2/tests/migration_resume.rs @@ -0,0 +1,97 @@ +use super::*; + +#[tokio::test] +async fn resume_recovers_empty_or_torn_first_migration_destination() { + for empty in [false, true] { + let root = tempfile::tempdir().unwrap(); + let workspace = root.path().to_path_buf(); + let id = crate::session::new_id(); + let expected = vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "legacy prompt"), + Item::text(ItemKind::Assistant, "legacy answer"), + ]; + let (legacy, scoped) = crate::session::test_support::interrupted_migration( + root.path(), + &id, + expected.clone(), + empty, + ); + let before = [ + std::fs::read(&legacy).unwrap(), + std::fs::read(&scoped).unwrap(), + ]; + crate::session::branch::validate_resume(root.path(), &id).unwrap(); + assert_eq!(std::fs::read(&legacy).unwrap(), before[0]); + assert_eq!(std::fs::read(&scoped).unwrap(), before[1]); + assert!(!legacy.with_extension("lock").exists()); + assert!(!scoped.with_extension("lock").exists()); + + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + ) + .unwrap(); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime, SessionRegistry::new()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let updates = Arc::new(Mutex::new(Vec::::new())); + let received = updates.clone(); + let resumed_id = id.clone(); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + async move |update: wire::UpdateSessionNotification, _cx| { + received.lock().unwrap().push(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("migration-resume-test", "0"), + )) + .block_task() + .await?; + cx.send_request( + wire::ResumeSessionRequest::new(resumed_id.clone(), workspace) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ) + .block_task() + .await?; + // This second request requires the recovered real actor to be + // live and settled, not merely a successful preflight response. + let listed = cx + .send_request(ListPromptBranchesRequest { + session_id: wire::SessionId::new(resumed_id.clone()), + }) + .block_task() + .await?; + assert_eq!(listed.boundaries.len(), 1); + cx.send_request(wire::CloseSessionRequest::new(resumed_id)) + .block_task() + .await?; + Ok(()) + }); + let result = timeout(Duration::from_secs(5), client).await; + server.abort(); + let _ = server.await; + result.expect("resume client timed out").unwrap(); + assert_eq!(crate::session::load(root.path(), &id).unwrap(), expected); + assert!(std::fs::read(&scoped).unwrap().ends_with(b"\n")); + let updates = updates.lock().unwrap(); + assert!(updates.iter().any(|update| matches!(&update.update, + wire::SessionUpdate::UserMessage(message) + if matches!(&message.content, MaybeUndefined::Value(content) + if content == &vec![wire::ContentBlock::Text(wire::TextContent::new("legacy prompt"))])))); + assert!(!updates.iter().any(|update| matches!( + update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ))); + } +} diff --git a/src/protocols/acp/v2/tests/review_tests.rs b/src/protocols/acp/v2/tests/review_tests.rs new file mode 100644 index 0000000..d7a0318 --- /dev/null +++ b/src/protocols/acp/v2/tests/review_tests.rs @@ -0,0 +1,582 @@ +mod partial_delivery; + +use super::*; + +/// The provider boundary records actual requests, not driver implementation work. +struct ContextAdapter { + turns: Arc, + requests: Arc>>>, +} + +struct ContextSession { + inner: TestSession, + requests: Arc>>>, +} + +#[async_trait] +impl ModelAdapter for ContextAdapter { + type Session = ContextSession; + + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(ContextSession { + inner: TestAdapter { + outcome: TestOutcome::Content, + turns: self.turns.clone(), + interrupt: None, + } + .start_session(config) + .await?, + requests: self.requests.clone(), + }) + } +} + +#[async_trait] +impl ModelSession for ContextSession { + type Turn = TestTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + self.requests + .lock() + .unwrap() + .push(request.transcript.clone()); + self.inner.begin_turn(request, cancellation).await + } +} + +struct InjectionWire { + client: agent_client_protocol::Channel, + server: tokio::task::JoinHandle>, + receipts: mpsc::UnboundedReceiver, +} + +impl Drop for InjectionWire { + fn drop(&mut self) { + self.server.abort(); + } +} + +impl InjectionWire { + async fn connect( + integration: Arc, + work: Arc>, + ) -> (Self, V2ConnectionTo) { + let (mut client, agent) = agent_client_protocol::Channel::duplex(); + let (connections, mut connection) = mpsc::unbounded_channel(); + let (accepted, receipts) = mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + agent_client_protocol::Agent + .v2() + .on_receive_request( + async move |request: wire::InitializeRequest, responder, cx| { + connections.send(cx.clone()).unwrap(); + responder.respond( + wire::InitializeResponse::new( + request.protocol_version, + wire::Implementation::new("review-injection-test", "0"), + ) + .capabilities(agentkit_acp::v2::agent_capabilities()), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: wire::InjectSessionRequest, responder, cx| { + let integration = integration.clone(); + let accepted = accepted.clone(); + let work = work.clone(); + cx.spawn(async move { + let Some(reserved) = integration + .reserve_inject_request(request, responder) + .await? + else { + return Ok(()); + }; + let id = reserved.response().message_id; + work.lock().unwrap().pending.insert(id.clone()); + let mut pending = TrackedInjection { + work, + id, + retained: false, + }; + let receipt = reserved.respond_tracked()?.expect("tracked acceptance"); + pending.retained = true; + accepted.send(receipt).ok(); + Ok(()) + })?; + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(agent) + .await + }); + send_wire( + &client, + "initialize", + 1, + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("review-injection-client", "0"), + )) + .unwrap(), + ); + assert!(receive_wire(&mut client).await.get("result").is_some()); + let connection = connection.recv().await.unwrap(); + ( + Self { + client, + server, + receipts, + }, + connection, + ) + } + + async fn inject( + &mut self, + session_id: &wire::SessionId, + request_id: i64, + text: &str, + ) -> agentkit_acp::v2::AcpInjectAcceptance { + send_wire( + &self.client, + "session/inject", + request_id, + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new(text))], + )) + .unwrap(), + ); + let response = receive_wire(&mut self.client).await; + assert!( + response.get("result").is_some(), + "injection failed: {response}" + ); + self.receipts.recv().await.unwrap() + } + + async fn update(&mut self) -> wire::UpdateSessionNotification { + let message = receive_wire(&mut self.client).await; + assert_eq!(message["method"], "session/update", "{message}"); + serde_json::from_value(message["params"].clone()).unwrap() + } +} + +#[derive(Clone, Copy, Debug)] +enum Entry { + Prompt, + Autonomous, + InitialBranch, + SetConfig, +} + +fn text_count(transcript: &[Item], expected: &str) -> usize { + transcript + .iter() + .filter(|item| { + item.kind == ItemKind::User + && item + .parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text == expected)) + }) + .count() +} + +async fn snapshot_actor(commands: &mpsc::Sender) -> SessionSnapshot { + let (reply, response) = oneshot::channel(); + commands.send(Command::Snapshot { reply }).await.unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .expect("cancelled delivery must retain the same actor") + .unwrap() +} + +async fn queue_prompt( + commands: &mpsc::Sender, + session_id: &wire::SessionId, + handle: &AcpSessionHandle, + text: &str, +) -> oneshot::Sender<()> { + let (reply, response) = oneshot::channel(); + commands + .send(Command::Prompt(PromptCommand { + request: wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new(text))], + ), + cancellation_generation: handle.cancellation_handle().generation(), + reply, + })) + .await + .unwrap(); + timeout(Duration::from_secs(5), response) + .await + .unwrap() + .unwrap() + .unwrap() +} + +async fn cancelled_delivery_keeps_actor(entry: Entry) { + const FIRST: &str = "/compact cancelled steering"; + const SECOND: &str = "second unactivated steering"; + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let durable_id = crate::session::new_id(); + let session_id = wire::SessionId::new(durable_id.clone()); + let loop_id = SessionId::new(durable_id.clone()); + let opened = crate::session::open( + root.path(), + &durable_id, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "older content ".repeat(20_000)), + Item::text(ItemKind::Assistant, "recent response") + .with_usage(Usage::new(agentkit_core::TokenUsage::new(100, 0))), + ], + ) + .unwrap(); + let integration = Arc::new(AcpIntegration::default()); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let (mut wire_client, connection) = + InjectionWire::connect(integration.clone(), work.clone()).await; + let sink = ResponseReplacementSink::new(ConnectionSink(connection, work.clone())); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + handle.prepare_injection_turn(); + handle.start_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let first = wire_client.inject(&session_id, 2, FIRST).await; + let first_id = first.message_id().clone(); + first.activate_after_response().await.unwrap(); + let second = wire_client.inject(&session_id, 3, SECOND).await; + let second_id = second.message_id().clone(); + assert_eq!(work.lock().unwrap().pending.len(), 2); + + let selection = SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let summaries = Arc::new(Mutex::new(Vec::new())); + let compactor = crate::compaction::automatic( + SwitchSummaryAdapter { + selection: selection.clone(), + seen: summaries.clone(), + outcome: TestOutcome::Content, + interrupt: None, + }, + Default::default(), + Some(opened.observer.clone()), + loop_id.clone(), + ) + .unwrap(); + let turns = Arc::new(AtomicU64::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let settlement = crate::runtime::InputSettlement::default(); + let initial_branch = matches!(entry, Entry::InitialBranch); + let driver = Agent::builder() + .mutator(settlement.clone()) + .mutator(compactor) + .model(ContextAdapter { + turns: turns.clone(), + requests: requests.clone(), + }) + .task_manager(manager) + .observer(ResponseReplacementObserver::new( + (*integration).clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .transcript_observer(settlement.observer(opened.observer)) + .transcript(opened.transcript) + .input(if initial_branch { + vec![Item::text(ItemKind::User, "initial branch prompt")] + } else { + vec![] + }) + .cancellation(handle.cancellation_handle()) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let busy = Arc::new(AtomicBool::new(initial_branch)); + let (commands, receiver) = mpsc::channel(8); + let mcp = crate::tools::mcp::empty(); + let events = mcp.subscribe(durable_id.clone()); + if matches!(entry, Entry::Autonomous) { + mcp.publish( + &durable_id, + crate::tools::mcp::McpEvent { + message: "autonomous trigger".into(), + }, + ); + } + let actor = tokio::spawn(session_actor(SessionActor { + initial_generation: initial_branch.then_some(generation), + admission_released: Arc::new(Notify::new()), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver: settlement.wrap(driver), + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: selection.clone(), + catalog: vec![crate::provider::ModelGroup { + provider: ProviderKind::OpenAiSubscription, + models: vec!["gpt-5.4-mini".into()], + context_windows: [("gpt-5.4-mini".into(), 150)].into_iter().collect(), + }], + commands: receiver, + mcp_events: events, + })); + let mut config_response = None; + match entry { + Entry::Prompt => { + claim_prompt(&busy).unwrap(); + queue_prompt(&commands, &session_id, &handle, "original prompt") + .await + .send(()) + .unwrap(); + } + Entry::SetConfig => { + let mut request = wire::SetSessionConfigOptionRequest::new( + session_id.clone(), + super::super::super::MODEL_CONFIG_ID, + "openai-subscription:gpt-5.4-mini", + ); + let (reply, response) = oneshot::channel(); + commands + .send(Command::SetConfig { + request: request.clone(), + reply, + cancellation_generation: generation, + }) + .await + .unwrap(); + let warning = response.await.unwrap().unwrap_err(); + let warning: model_switch::Warning = + serde_json::from_value(warning.data.unwrap()[model_switch::META].clone()).unwrap(); + assert!(summaries.lock().unwrap().is_empty()); + request.meta = Some(serde_json::Map::from_iter([( + model_switch::META.into(), + serde_json::to_value(model_switch::Confirmation { + token: warning.token, + action: model_switch::Decision::Compact, + }) + .unwrap(), + )])); + let (reply, response) = oneshot::channel(); + commands + .send(Command::SetConfig { + request, + reply, + cancellation_generation: generation, + }) + .await + .unwrap(); + config_response = Some(response); + } + Entry::Autonomous | Entry::InitialBranch => {} + } + + let mut updates = Vec::new(); + loop { + let update = wire_client.update().await; + let delivered = matches!(&update.update, wire::SessionUpdate::UserMessage(message) + if message.message_id == first_id); + updates.push(update); + if delivered { + break; + } + } + assert!(!work.lock().unwrap().pending.contains(&first_id)); + assert!(work.lock().unwrap().pending.contains(&second_id)); + let calls_before_cancel = turns.load(Ordering::Relaxed); + let summaries_before_cancel = summaries.lock().unwrap().len(); + assert_eq!( + summaries_before_cancel, + usize::from(matches!(entry, Entry::SetConfig)) + ); + assert_eq!( + text_count( + &crate::session::load(root.path(), &durable_id).unwrap(), + FIRST + ), + 0, + "the delivered notification precedes persistence while the second receipt holds the boundary" + ); + handle.interrupt(); + loop { + let update = wire_client.update().await; + let idle = matches!( + &update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ); + updates.push(update); + if idle { + break; + } + } + let states = updates + .iter() + .filter(|update| matches!(update.update, wire::SessionUpdate::StateUpdate(_))) + .cloned() + .collect::>(); + assert_running_then_idle(&states, wire::StopReason::Cancelled); + assert!(!updates.iter().any( + |update| matches!(&update.update, wire::SessionUpdate::UserMessage(message) + if message.message_id == second_id) + )); + if let Some(response) = config_response { + assert!(response.await.unwrap().is_err()); + assert_eq!(selection.selection().unwrap().model, "gpt-5.4"); + } + assert_eq!( + text_count( + &crate::session::load(root.path(), &durable_id).unwrap(), + FIRST + ), + 1, + "cancellation must persist already acknowledged steering" + ); + let settled = snapshot_actor(&commands).await; + assert_eq!(text_count(&settled.canonical_transcript, FIRST), 1); + assert_eq!(text_count(&settled.canonical_transcript, SECOND), 0); + let persisted = crate::session::load(root.path(), &durable_id).unwrap(); + assert_eq!(persisted, settled.canonical_transcript); + assert_eq!(turns.load(Ordering::Relaxed), calls_before_cancel); + assert_eq!(summaries.lock().unwrap().len(), summaries_before_cancel); + assert!(work.lock().unwrap().pending.contains(&second_id)); + + // A real queued Prompt response gate proves the same actor remains usable + // without accidentally activating the second receipt or executing old work. + claim_prompt(&busy).unwrap(); + handle.prepare_injection_turn(); + let start = queue_prompt(&commands, &session_id, &handle, "fresh generation prompt").await; + assert_eq!(turns.load(Ordering::Relaxed), calls_before_cancel); + assert!(work.lock().unwrap().pending.contains(&second_id)); + second.activate_after_response().await.unwrap(); + start.send(()).unwrap(); + let mut fresh_updates = Vec::new(); + loop { + let update = wire_client.update().await; + let idle = matches!( + &update.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ); + fresh_updates.push(update); + if idle { + break; + } + } + let states = fresh_updates + .iter() + .filter(|update| matches!(update.update, wire::SessionUpdate::StateUpdate(_))) + .cloned() + .collect::>(); + assert_running_then_idle(&states, wire::StopReason::EndTurn); + assert_eq!( + fresh_updates + .iter() + .filter(|update| matches!(&update.update, + wire::SessionUpdate::UserMessage(message) if message.message_id == second_id)) + .count(), + 1 + ); + let final_snapshot = snapshot_actor(&commands).await; + assert_eq!(text_count(&final_snapshot.canonical_transcript, FIRST), 1); + assert_eq!(text_count(&final_snapshot.canonical_transcript, SECOND), 1); + assert_eq!(summaries.lock().unwrap().len(), summaries_before_cancel); + assert!(work.lock().unwrap().pending.is_empty()); + assert!(requests.lock().unwrap().iter().any(|transcript| text_count( + transcript, + "fresh generation prompt" + ) == 1 + && text_count(transcript, FIRST) == 1)); + assert_eq!( + crate::session::load(root.path(), &durable_id).unwrap(), + final_snapshot.canonical_transcript + ); + let (reply, response) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + response.await.unwrap(); + actor.await.unwrap(); +} + +#[tokio::test] +async fn delivered_steering_cancellation_retains_prompt_actor() { + timeout( + Duration::from_secs(20), + cancelled_delivery_keeps_actor(Entry::Prompt), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn delivered_steering_cancellation_retains_autonomous_actor() { + timeout( + Duration::from_secs(20), + cancelled_delivery_keeps_actor(Entry::Autonomous), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn delivered_steering_cancellation_retains_initial_branch_actor() { + timeout( + Duration::from_secs(20), + cancelled_delivery_keeps_actor(Entry::InitialBranch), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn delivered_steering_cancellation_retains_compaction_actor() { + timeout( + Duration::from_secs(20), + cancelled_delivery_keeps_actor(Entry::SetConfig), + ) + .await + .unwrap(); +} diff --git a/src/protocols/acp/v2/tests/review_tests/partial_delivery.rs b/src/protocols/acp/v2/tests/review_tests/partial_delivery.rs new file mode 100644 index 0000000..3c2859e --- /dev/null +++ b/src/protocols/acp/v2/tests/review_tests/partial_delivery.rs @@ -0,0 +1,426 @@ +use super::*; +use futures_util::StreamExt; +use std::sync::OnceLock; + +/// Fail at the actual transport acknowledgement boundary, not at a driver hook. +/// The immutable fault is installed before the actor starts. Only the external +/// send-attempt counter changes here; normal ConnectionSink tracking is untouched. +/// Retain only the cancellation controller, never the session that owns this sink. +#[derive(Clone)] +struct PartialDeliverySink { + inner: ConnectionSink, + fault: Arc>, + cancellation: CancellationController, + failures: Arc, + fail_after_observation: Arc, +} + +#[async_trait] +impl AcpSessionUpdateSink for PartialDeliverySink { + fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError> { + self.inner.update(notification) + } + + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + if let Some(failed_id) = self.fault.get() + && matches!(¬ification.update, wire::SessionUpdate::UserMessage(message) + if &message.message_id == failed_id) + { + self.failures.fetch_add(1, Ordering::Relaxed); + // B has already entered driver pending_input, but neither the real + // transport nor ConnectionSink's delivered bookkeeping accepts it. + self.fail_after_observation.notified().await; + self.cancellation.interrupt(); + return Err(AcpRuntimeError::Sdk( + "injected B notification failure".into(), + )); + } + self.inner.update_acknowledged(notification).await + } + + async fn flush(&self) -> Result<(), AcpRuntimeError> { + self.inner.flush().await + } +} + +async fn inject_content( + client: &mut InjectionWire, + session_id: &wire::SessionId, + request_id: i64, + content: Vec, +) -> wire::MessageId { + send_wire( + &client.client, + "session/inject", + request_id, + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + content, + )) + .unwrap(), + ); + let response = receive_wire(&mut client.client).await; + assert!(response.get("result").is_some(), "{response}"); + let receipt = client.receipts.recv().await.unwrap(); + let id = receipt.message_id().clone(); + receipt.activate_after_response().await.unwrap(); + id +} + +fn content_occurrences(transcript: &[Item], content: &[wire::ContentBlock], expected: usize) { + for item in agentkit_acp::v2::content_blocks_to_items(content).unwrap() { + assert_eq!( + transcript + .iter() + .filter(|actual| actual.kind == item.kind && actual.parts == item.parts) + .count(), + expected, + "wrong durable multiplicity for {:?}", + item.parts, + ); + } +} + +async fn partial_delivery_retires_actor(multi_item: bool, identical: bool) { + let first_text = if identical { + "identical steering" + } else { + "/compact acknowledged steering" + }; + let second_text = if identical { + first_text + } else { + "unacknowledged steering B" + }; + let content = |text: &str, resource: &str| { + let mut content = vec![wire::ContentBlock::Text(wire::TextContent::new(text))]; + if multi_item { + content.push(wire::ContentBlock::ResourceLink(wire::ResourceLink::new( + resource, + format!("file:///{resource}"), + ))); + } + content + }; + let first_content = content(first_text, "acknowledged-context-A"); + let second_content = content(second_text, "unacknowledged-context-B"); + assert_eq!( + agentkit_acp::v2::content_blocks_to_items(&first_content) + .unwrap() + .len(), + if multi_item { 2 } else { 1 }, + "the multi-item case must exercise a prefix of items, not notifications", + ); + let root = tempfile::tempdir().unwrap(); + let durable_id = crate::session::new_id(); + let session_id = wire::SessionId::new(durable_id.clone()); + let integration = Arc::new(AcpIntegration::default()); + let turns = Arc::new(AtomicU64::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let summaries = Arc::new(Mutex::new(Vec::new())); + + // Attempt 1 is a new attachment and disk reload, not an automatic retry on + // the failed driver. Reuse the integration and session ID to test unbinding. + for reconnect in [false, true] { + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let loop_id = SessionId::new(durable_id.clone()); + let opened = crate::session::open( + root.path(), + &durable_id, + reconnect, + false, + if reconnect { + vec![] + } else { + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "older content ".repeat(20_000)), + Item::text(ItemKind::Assistant, "recent response") + .with_usage(Usage::new(agentkit_core::TokenUsage::new(100, 0))), + ] + }, + ) + .unwrap(); + if reconnect { + content_occurrences(&opened.transcript, &first_content, 1); + content_occurrences(&opened.transcript, &second_content, usize::from(identical)); + } + let work = Arc::new(Mutex::new(InjectionWork::default())); + let (mut client, connection) = + InjectionWire::connect(integration.clone(), work.clone()).await; + let fault = Arc::new(OnceLock::new()); + let cancellation = CancellationController::new(); + let failures = Arc::new(AtomicU64::new(0)); + let fail_after_observation = Arc::new(Notify::new()); + let sink = ResponseReplacementSink::new(PartialDeliverySink { + inner: ConnectionSink(connection, work.clone()), + fault: fault.clone(), + cancellation: cancellation.clone(), + failures: failures.clone(), + fail_after_observation: fail_after_observation.clone(), + }); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session( + AcpSessionBinding::new(session_id.clone(), loop_id.clone(), sink.clone()) + .cancellation(cancellation.clone()), + ) + .unwrap(); + handle.prepare_injection_turn(); + handle.start_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let first_id = if reconnect { + None + } else { + Some(inject_content(&mut client, &session_id, 2, first_content.clone()).await) + }; + let second_id = inject_content(&mut client, &session_id, 3, second_content.clone()).await; + assert_eq!( + work.lock().unwrap().pending.len(), + if reconnect { 1 } else { 2 } + ); + if !reconnect { + assert!(fault.set(second_id.clone()).is_ok()); + } + + let selection = SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let compactor = crate::compaction::automatic( + SwitchSummaryAdapter { + selection: selection.clone(), + seen: summaries.clone(), + outcome: TestOutcome::Content, + interrupt: None, + }, + Default::default(), + Some(opened.observer.clone()), + loop_id.clone(), + ) + .unwrap(); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let settlement = crate::runtime::InputSettlement::default(); + let driver = Agent::builder() + .mutator(settlement.clone()) + .mutator(compactor) + .model(ContextAdapter { + turns: turns.clone(), + requests: requests.clone(), + }) + .task_manager(manager) + .observer(ResponseReplacementObserver::new( + (*integration).clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .transcript_observer(settlement.observer(opened.observer)) + .transcript(opened.transcript) + .cancellation(handle.cancellation_handle()) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let busy = Arc::new(AtomicBool::new(false)); + let (commands, receiver) = mpsc::channel(8); + let mcp = crate::tools::mcp::empty(); + let actor = tokio::spawn(session_actor(SessionActor { + initial_generation: None, + admission_released: Arc::new(Notify::new()), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration: integration.clone(), + session_id: session_id.clone(), + }, + sink, + activity, + driver: settlement.wrap(driver), + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: selection, + catalog: vec![], + commands: receiver, + mcp_events: mcp.subscribe(durable_id.clone()), + })); + claim_prompt(&busy).unwrap(); + queue_prompt(&commands, &session_id, &handle, "explicit prompt") + .await + .send(()) + .unwrap(); + + if !reconnect { + // Observe A through the real Channel while B's transport send is + // suspended. Both acceptance receipts were activated before run. + loop { + let update = client.update().await; + if matches!(&update.update, wire::SessionUpdate::UserMessage(message) + if Some(&message.message_id) == first_id.as_ref()) + { + break; + } + } + let calls_before_cancel = turns.load(Ordering::Relaxed); + let requests_before_cancel = requests.lock().unwrap().len(); + let summaries_before_cancel = summaries.lock().unwrap().len(); + assert_eq!(summaries_before_cancel, 0); + fail_after_observation.notify_one(); + + // An Idle(Cancelled) notification is not proof that this attachment + // can be reused. The command receiver and BindingGuard must retire. + timeout(Duration::from_secs(5), actor) + .await + .expect("partial transport failure must retire, not idle or retry") + .unwrap(); + assert!(commands.is_closed()); + let (reply, _) = oneshot::channel(); + assert!(commands.send(Command::Snapshot { reply }).await.is_err()); + assert!(matches!( + integration.flush_session_updates(&session_id).await, + Err(AcpRuntimeError::SessionNotFound(_)) + )); + assert!(handle.cancellation_handle().is_cancelled_since(generation)); + assert_eq!(failures.load(Ordering::Relaxed), 1, "B must not auto-retry"); + assert_eq!(turns.load(Ordering::Relaxed), calls_before_cancel); + assert_eq!(requests.lock().unwrap().len(), requests_before_cancel); + assert!( + summaries.lock().unwrap().len() == summaries_before_cancel, + "/compact must not execute" + ); + let pending = work.lock().unwrap().pending.clone(); + assert!(!pending.contains(first_id.as_ref().unwrap())); + assert_eq!(pending.len(), 1); + assert!( + pending.contains(&second_id), + "failed B must not be marked delivered" + ); + let persisted = crate::session::load(root.path(), &durable_id).unwrap(); + content_occurrences(&persisted, &first_content, 1); + content_occurrences(&persisted, &second_content, usize::from(identical)); + let disk = std::fs::read_to_string(crate::session::transcript_path_for_test( + root.path(), + &durable_id, + )) + .unwrap(); + assert_eq!(disk.matches(first_text).count(), 1); + if !identical { + assert!( + !disk.contains(second_text), + "B must never be appended, even transiently" + ); + } + if multi_item { + assert!(!disk.contains("unacknowledged-context-B")); + } + // Drain everything already emitted after joining the actor, without + // waiting for Idle (retirement is required even if Idle cannot send). + let mut delivered = vec![first_id.clone().unwrap()]; + while let Some(frame) = client.client.rx.next().now_or_never().flatten() { + let agent_client_protocol::TransportFrame::Single(message) = frame else { + panic!("unexpected batch frame"); + }; + let value = serde_json::to_value(message).unwrap(); + if value["method"] == "session/update" { + let update: wire::UpdateSessionNotification = + serde_json::from_value(value["params"].clone()).unwrap(); + if let wire::SessionUpdate::UserMessage(message) = update.update + && (Some(&message.message_id) == first_id.as_ref() + || message.message_id == second_id) + { + delivered.push(message.message_id); + } + } + } + assert_eq!(delivered, vec![first_id.unwrap()]); + } else { + let mut delivered = Vec::new(); + loop { + let update = client.update().await; + match update.update { + wire::SessionUpdate::UserMessage(message) + if message.message_id == second_id => + { + delivered.push(message.message_id); + } + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) => { + assert_eq!(idle.stop_reason, Some(wire::StopReason::EndTurn)); + break; + } + _ => {} + } + } + assert_eq!(delivered, vec![second_id]); + assert!(work.lock().unwrap().pending.is_empty()); + assert_eq!(failures.load(Ordering::Relaxed), 0); + let snapshot = snapshot_actor(&commands).await; + content_occurrences( + &snapshot.canonical_transcript, + &first_content, + 1 + usize::from(identical), + ); + content_occurrences( + &snapshot.canonical_transcript, + &second_content, + 1 + usize::from(identical), + ); + assert_eq!( + crate::session::load(root.path(), &durable_id).unwrap(), + snapshot.canonical_transcript, + ); + let (reply, response) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + response.await.unwrap(); + actor.await.unwrap(); + } + } +} + +#[tokio::test] +async fn partial_transport_error_and_cancel_retire_actor() { + timeout( + Duration::from_secs(20), + partial_delivery_retires_actor(false, false), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn partial_transport_error_preserves_all_items_of_acknowledged_steer() { + timeout( + Duration::from_secs(20), + partial_delivery_retires_actor(true, false), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn partial_transport_error_distinguishes_identical_steers_by_position() { + timeout( + Duration::from_secs(20), + partial_delivery_retires_actor(false, true), + ) + .await + .unwrap(); +} diff --git a/src/runtime.rs b/src/runtime.rs index 72690c1..0abb61d 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -44,6 +44,9 @@ use crate::{ }, }; +mod input_settlement; +pub(crate) use input_settlement::{InputSettlement, InputSettlingDriver}; + #[cfg(test)] mod test_support { use super::*; @@ -197,6 +200,49 @@ impl SessionSelection { } } +fn acp_branch_selection( + transcript: &[Item], +) -> Result)>, AcpRuntimeError> { + crate::session::branch::BranchMetadata::read(transcript) + .and_then(|metadata| { + metadata + .map(|metadata| metadata.request.selection.resolve()) + .transpose() + }) + .map_err(AcpRuntimeError::Loop) +} + +fn ordinary_acp_fork_transcript(mut transcript: Vec) -> Result, AcpRuntimeError> { + crate::session::branch::strip_for_plain_fork(&mut transcript).map_err(AcpRuntimeError::Loop)?; + Ok(transcript) +} + +/// A branch's first input is already part of its durable completion snapshot. +/// The loop must receive it as pending input to activate, but must not append it +/// to disk a second time. All other transcript events retain normal persistence. +struct AcpTranscriptObserver { + observer: crate::session::SessionObserver, + committed_input: Mutex>, +} + +impl agentkit_loop::TranscriptObserver for AcpTranscriptObserver { + fn on_transcript_event(&self, event: agentkit_loop::TranscriptEvent<'_>) { + let committed = self + .committed_input + .lock() + .expect("ACP committed input guard poisoned") + .take(); + if let Some(committed) = committed { + assert_eq!( + &committed, event.item, + "branch activation must append its committed prompt first" + ); + } else { + self.observer.on_transcript_event(event); + } + } +} + pub(crate) struct AcpForkState { pub transcript: Vec, pub selection: ModelSelection, @@ -332,7 +378,7 @@ impl Drop for SessionClaim { } pub(crate) struct AcpDriver { - pub driver: LoopDriver, + pub driver: InputSettlingDriver, pub skills: Vec, pub tasks: TaskManagerHandle, pub background_jobs: BackgroundJobs, @@ -1089,6 +1135,26 @@ impl Runtime { .clone() .ok_or_else(|| "persistent run requires a configured session".to_string())?; let session_id = request.id.clone(); + let branch = if request.resume { + crate::session::branch::validate_resume(&self.root, &request.id)? + } else { + None + }; + // Noninteractive reloads must honor the same persisted child selection + // as ACP, before starting an adapter or changing the transcript. + let adapter = if let Some(metadata) = branch { + let (selection, reasoning) = metadata.request.selection.resolve()?; + SelectableAdapter::new_with_credentials_effort_and_openrouter_key( + selection.provider, + selection.model, + self.credential_storage.clone(), + reasoning, + self.openrouter_api_key.clone(), + )? + } else { + self.adapter.clone() + }; + if self.plugin_runtime.is_some() { self.mcp.refresh().await.map_err(|error| { record_runtime_failure( @@ -1127,7 +1193,7 @@ impl Runtime { let pending_creation = (!request.resume).then(|| opened.observer.clone()); let skills = self.fresh_skills(); let compactor = crate::compaction::automatic( - self.adapter.clone(), + adapter.clone(), self.agentkit_telemetry(), Some(opened.observer.clone()), format!("compaction-{}", crate::session::new_id()), @@ -1142,7 +1208,7 @@ impl Runtime { })?; let subagents = self.subagents.fresh(); let agent = Agent::builder() - .model(self.adapter.clone()) + .model(adapter) .telemetry(self.agentkit_telemetry()) .add_tool_source(self.compose_with_jobs( 0, @@ -1336,6 +1402,40 @@ impl Runtime { claim: &mut SessionClaim, forked: Option, ) -> Result + where + I: LoopObserver + Clone + 'static, + { + self.start_acp_driver_with_persistence(context, claim, forked, false) + .await + } + + /// Commit a prepared prompt checkout before any adapter or model startup. + /// Unlike ordinary forks, a committed child survives a failed start/response. + pub(crate) async fn start_acp_branch_driver_with_initial( + self: &Arc, + context: AcpDriverContext, + claim: &mut SessionClaim, + forked: AcpForkState, + ) -> Result + where + I: LoopObserver + Clone + 'static, + { + if !claim.is_fork() { + return Err(AcpRuntimeError::Loop( + "a prompt checkout requires a fork claim".into(), + )); + } + self.start_acp_driver_with_persistence(context, claim, Some(forked), true) + .await + } + + async fn start_acp_driver_with_persistence( + self: &Arc, + context: AcpDriverContext, + claim: &mut SessionClaim, + forked: Option, + commit_branch: bool, + ) -> Result where I: LoopObserver + Clone + 'static, { @@ -1364,8 +1464,14 @@ impl Runtime { "a forked transcript requires a new session identity".into(), )); } - transcript + if commit_branch { + transcript + } else { + ordinary_acp_fork_transcript(transcript)? + } } else if request.resume { + crate::session::branch::validate_resume(&self.root, &request.id) + .map_err(AcpRuntimeError::Loop)?; vec![Item::text( ItemKind::System, self.system_prompt(self.base_depth), @@ -1386,6 +1492,23 @@ impl Runtime { if is_fork || !request.resume { claim.guard_uncommitted_transcript(&opened.observer); } + // Validate the opened payload as well, then restore a resumed checkout's + // captured selection. Ordinary forks already cleared validated ancestry. + let persisted = acp_branch_selection(&opened.transcript)?; + let selected = if request.resume || commit_branch { + persisted.or(selected) + } else { + // Ordinary forks inherit the parent's *current* model, not the + // model recorded when an ancestor checkout was created. + selected + }; + if commit_branch { + crate::session::branch::commit(&opened.observer, &opened.transcript) + .map_err(AcpRuntimeError::Loop)?; + // commit disarms the observer's creation guard after the disk + // barrier. The claim retains its observer, but dropping either + // owner after an adapter/start failure cannot delete this child. + } // 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(|| { @@ -1420,10 +1543,23 @@ impl Runtime { let tasks = task_manager.handle(); let background_jobs = BackgroundJobs::default(); let canonical_transcript = opened.transcript.clone(); + let mut transcript = opened.transcript; + let committed_input = if commit_branch { + // branch::commit already checked the complete prefix + user shape. + transcript.pop() + } else { + None + }; + let input = committed_input.iter().cloned().collect(); + let observer = AcpTranscriptObserver { + observer: opened.observer, + committed_input: Mutex::new(committed_input), + }; let mut session_config = SessionConfig::new(session_id.clone()).without_cache(); if context.response_attempt_replacement { session_config = session_config.with_response_attempt_supersession(); } + let input_settlement = InputSettlement::default(); let driver = Agent::builder() .model(adapter.clone()) .telemetry(self.agentkit_telemetry()) @@ -1434,10 +1570,13 @@ impl Runtime { skills, )) .task_manager(task_manager) + // Settlement must fence every mutator, including compaction. + .mutator(input_settlement.clone()) .mutator(compactor) .observer(context.integration.as_ref().clone()) - .transcript_observer(opened.observer) - .transcript(opened.transcript) + .transcript_observer(input_settlement.observer(observer)) + .transcript(transcript) + .input(input) .cancellation(context.cancellation) .build() .map_err(|error| AcpRuntimeError::Loop(error.to_string()))? @@ -1445,7 +1584,7 @@ impl Runtime { .await .map_err(|error| AcpRuntimeError::Loop(error.to_string()))?; let driver = AcpDriver { - driver, + driver: input_settlement.wrap(driver), skills: skill_catalog, tasks, background_jobs, @@ -2529,3 +2668,514 @@ impl Drop for StorageCancellationBridge { self.0.abort(); } } + +#[cfg(test)] +mod branch_startup_tests { + use super::*; + use crate::session::branch::{self, Boundary, CapturedSelection, SubmittedRequest}; + + fn prepared() -> Vec { + let prefix = vec![Item::text(ItemKind::System, "system")]; + branch::prepare( + prefix.clone(), + "parent".into(), + Boundary::new(0, &prefix).unwrap(), + "checkout".into(), + SubmittedRequest { + id: "request".into(), + selection: CapturedSelection::new( + &ModelSelection::new(ProviderKind::OpenRouter, "openai/gpt-5.4"), + Some(ReasoningEffort::High), + ), + }, + Item::text(ItemKind::User, "edited prompt"), + ) + .unwrap() + } + + #[test] + fn branch_resume_selection_is_read_without_mutation() { + let transcript = prepared(); + let before = transcript.clone(); + assert_eq!( + acp_branch_selection(&transcript).unwrap(), + Some(( + ModelSelection::new(ProviderKind::OpenRouter, "openai/gpt-5.4"), + Some(ReasoningEffort::High), + )) + ); + assert_eq!(transcript, before); + assert!( + acp_branch_selection(&[Item::text(ItemKind::System, "legacy")]) + .unwrap() + .is_none() + ); + } + + #[test] + fn branch_resume_rejects_invalid_payload_without_defaults() { + let original = prepared(); + let payload = original[0].metadata[branch::METADATA_KEY].clone(); + let mut invalid = vec![Value::Null, json!({ "version": 999 })]; + for (field, value) in [ + ("provider", json!("unknown")), + ("model", json!("")), + ("reasoning", json!("extreme")), + ] { + let mut changed = payload.clone(); + changed["request"]["selection"][field] = value; + invalid.push(changed); + } + let mut unknown = payload; + unknown["unrecognized"] = json!(true); + invalid.push(unknown); + for value in invalid { + let mut transcript = original.clone(); + transcript[0] + .metadata + .insert(branch::METADATA_KEY.into(), value); + assert!(acp_branch_selection(&transcript).is_err()); + } + } + + #[derive(Clone)] + struct QuietObserver; + + impl LoopObserver for QuietObserver { + fn handle_event(&self, _: agentkit_loop::ObservedEvent) {} + } + + fn context(root: &Path) -> AcpDriverContext { + AcpDriverContext { + cwd: root.into(), + additional_directories: Vec::new(), + integration: Arc::new(QuietObserver), + cancellation: CancellationController::new().handle(), + response_attempt_replacement: false, + } + } + + #[tokio::test] + async fn committed_branch_survives_failed_start_and_restores_selection_on_resume() { + let root = tempfile::tempdir().unwrap(); + // Empty explicit key deterministically fails OpenRouter adapter startup + // without credential lookup/network. Runtime defaults differ from child. + let runtime = Runtime::new_with_provider_credentials_effort_and_openrouter_key( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + let mut claim = runtime.claim_session_fork().unwrap(); + let id = claim.id().to_owned(); + let committed = prepared(); + let error = runtime + .start_acp_branch_driver_with_initial( + context(root.path()), + &mut claim, + AcpForkState { + transcript: committed.clone(), + selection: ModelSelection::new(ProviderKind::OpenAiSubscription, "gpt-5.4"), + reasoning_effort: None, + parent_context: None, + }, + ) + .await + .err() + .expect("empty explicit key must fail startup"); + assert!( + error + .to_string() + .contains("--openrouter-api-key cannot be empty"), + "{error}" + ); + drop(claim); + let recovered = branch::lookup_committed(root.path(), &id, "checkout", "request") + .unwrap() + .expect("failed adapter startup must remain recoverable"); + assert_eq!(recovered.transcript, committed); + let before = branch::load_history(root.path(), &id).unwrap(); + let path = crate::session::transcript_path_for_test(root.path(), &id); + let complete_bytes = std::fs::read(&path).unwrap(); + let mut torn_bytes = complete_bytes.clone(); + torn_bytes.extend_from_slice(br#"{"schema_version":3,"item":{"kind":"assistant""#); + std::fs::write(&path, &torn_bytes).unwrap(); + assert_eq!( + branch::lookup_committed(root.path(), &id, "checkout", "request") + .unwrap() + .unwrap() + .transcript, + committed, + ); + assert_eq!(std::fs::read(&path).unwrap(), torn_bytes); + let mut resume = runtime.claim_session_load(&id).unwrap(); + let error = runtime + .start_acp_driver(context(root.path()), &mut resume) + .await + .err() + .expect("resume must use the persisted OpenRouter selection"); + assert!( + error + .to_string() + .contains("--openrouter-api-key cannot be empty"), + "{error}" + ); + drop(resume); + assert_eq!(branch::load_history(root.path(), &id).unwrap(), before); + assert_eq!(crate::session::load(root.path(), &id).unwrap(), committed); + assert_eq!( + std::fs::read(&path).unwrap(), + complete_bytes, + "ACP opener only truncates the torn tail; no prompt or generation rerun" + ); + } + + #[tokio::test] + async fn interrupted_first_migration_resumes_through_runtime_and_acp() { + for empty in [false, true] { + for acp in [false, true] { + let root = tempfile::tempdir().unwrap(); + let id = "migration-resume"; + let expected = vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "legacy prompt"), + ]; + let (legacy, scoped) = crate::session::test_support::interrupted_migration( + root.path(), + id, + expected.clone(), + empty, + ); + let before = [ + std::fs::read(&legacy).unwrap(), + std::fs::read(&scoped).unwrap(), + ]; + assert_eq!(branch::validate_resume(root.path(), id).unwrap(), None); + assert_eq!(std::fs::read(&legacy).unwrap(), before[0]); + assert_eq!(std::fs::read(&scoped).unwrap(), before[1]); + assert!(!legacy.with_extension("lock").exists()); + assert!(!scoped.with_extension("lock").exists()); + + let runtime = Runtime::with_session_provider_credentials_effort_and_openrouter_key( + root.path(), + "openai/gpt-5.4", + ProviderKind::OpenRouter, + SessionRequest { + id: id.into(), + resume: true, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + let error = if acp { + let mut claim = runtime.claim_session_load(id).unwrap(); + runtime + .start_acp_driver(context(root.path()), &mut claim) + .await + .err() + .expect("adapter must reject empty key") + .to_string() + } else { + runtime + .run_persistent("must not be appended".into()) + .await + .unwrap_err() + }; + assert!( + error.contains("--openrouter-api-key cannot be empty"), + "{error}" + ); + assert_eq!(crate::session::load(root.path(), id).unwrap(), expected); + } + } + } + + #[test] + fn interrupted_first_migration_preflight_rejects_invalid_authority_without_mutation() { + for case in [ + "malformed-destination", + "terminated-torn-destination", + "malformed-legacy", + "foreign-workspace", + "foreign-session", + "divergent-history", + "incomplete-checkout", + "missing-legacy", + ] { + let root = tempfile::tempdir().unwrap(); + let id = "invalid-migration"; + let transcript = if case == "incomplete-checkout" { + prepared()[..1].to_vec() + } else { + vec![Item::text(ItemKind::System, "system")] + }; + let (legacy, scoped) = crate::session::test_support::interrupted_migration( + root.path(), + id, + transcript, + false, + ); + match case { + "malformed-destination" => std::fs::write(&scoped, b"{}").unwrap(), + "terminated-torn-destination" => std::fs::write(&scoped, b"{\n").unwrap(), + "malformed-legacy" => std::fs::write(&legacy, b"{}\n").unwrap(), + "foreign-workspace" | "foreign-session" => { + let mut record: serde_json::Value = + serde_json::from_slice(&std::fs::read(&legacy).unwrap()).unwrap(); + if case == "foreign-workspace" { + record["workspace_root"] = json!(root.path().join("elsewhere")); + } else { + record["session_id"] = json!("unauthorized"); + } + std::fs::write(&legacy, serde_json::to_vec(&record).unwrap()).unwrap(); + } + "divergent-history" => { + let mut record: serde_json::Value = + serde_json::from_slice(&std::fs::read(&legacy).unwrap()).unwrap(); + record["replacement"] = json!([Item::text(ItemKind::System, "unrelated")]); + std::fs::write(&scoped, serde_json::to_vec(&record).unwrap()).unwrap(); + } + "missing-legacy" => std::fs::remove_file(&legacy).unwrap(), + "incomplete-checkout" => {} + _ => unreachable!(), + } + let legacy_before = std::fs::read(&legacy).ok(); + let scoped_before = std::fs::read(&scoped).unwrap(); + assert!(branch::validate_resume(root.path(), id).is_err(), "{case}"); + assert_eq!(std::fs::read(&legacy).ok(), legacy_before, "{case}"); + assert_eq!(std::fs::read(&scoped).unwrap(), scoped_before, "{case}"); + assert!(!legacy.with_extension("lock").exists(), "{case}"); + assert!(!scoped.with_extension("lock").exists(), "{case}"); + } + } + + #[tokio::test] + async fn noninteractive_branch_reload_validates_completion_and_restores_selection() { + let root = tempfile::tempdir().unwrap(); + for complete in [false, true] { + let id = if complete { + "complete-print" + } else { + "partial-print" + }; + let opened = + crate::session::open_uncommitted(root.path(), id, false, prepared()).unwrap(); + if complete { + branch::commit(&opened.observer, &opened.transcript).unwrap(); + } else { + opened.observer.commit_creation().unwrap(); // interrupted initial appends, no completion + } + drop(opened); + let path = crate::session::transcript_path_for_test(root.path(), id); + let complete_bytes = std::fs::read(&path).unwrap(); + let mut torn_bytes = complete_bytes.clone(); + torn_bytes.extend_from_slice(br#"{"schema_version":3,"item":{"kind":"tool""#); + std::fs::write(&path, &torn_bytes).unwrap(); + assert_eq!(branch::validate_resume(root.path(), id).is_ok(), complete); + assert!( + std::fs::read(&path).unwrap() == torn_bytes, + "resume preflight must not repair" + ); + let runtime = Runtime::with_session_provider_credentials_effort_and_openrouter_key( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + SessionRequest { + id: id.into(), + resume: true, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + let error = runtime + .run_persistent("must not be appended".into()) + .await + .unwrap_err(); + assert!( + error.contains(if complete { + "--openrouter-api-key cannot be empty" + } else { + "incomplete prompt checkout" + }), + "{error}" + ); + let expected = if complete { complete_bytes } else { torn_bytes }; + assert!( + std::fs::read(&path).unwrap() == expected, + "only a committed resume's locked opener may repair; no generation rerun" + ); + } + } + + #[tokio::test] + async fn partial_branch_resume_rejects_before_adapter_or_disk_repair() { + use agentkit_loop::{TranscriptEvent, TranscriptObserver}; + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_credentials_effort_and_openrouter_key( + root.path(), + "gpt-5.4", + ProviderKind::OpenRouter, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + for retained in [1, 2] { + let id = format!("partial-{retained}"); + let transcript = prepared(); + let opened = crate::session::open_uncommitted( + root.path(), + &id, + false, + transcript[..retained].to_vec(), + ) + .unwrap(); + if retained == 2 { + // If resume reached open, this unanswered call would cause a + // repair append. Read-only preflight must leave it untouched. + let mut call = Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(agentkit_core::ToolCallPart::new( + "unanswered", + "compose", + json!({}), + ))], + ); + call.created_at = transcript[1].created_at; + opened.observer.on_transcript_event(TranscriptEvent { + session_id: &agentkit_core::SessionId::new(&id), + item: &call, + }); + } + // Simulate a crash leaving initial appends, but no branch commit. + opened.observer.commit_creation().unwrap(); + drop(opened); + let before = branch::load_history(root.path(), &id).unwrap(); + let mut claim = runtime.claim_session_load(&id).unwrap(); + let error = runtime + .start_acp_driver(context(root.path()), &mut claim) + .await + .err() + .expect("initial appends cannot resume a checkout"); + assert!( + error.to_string().contains("incomplete prompt checkout"), + "{error}" + ); + assert!(error.to_string().contains(&id), "{error}"); + drop(claim); + assert_eq!(branch::load_history(root.path(), &id).unwrap(), before); + } + } + + #[tokio::test] + async fn ordinary_branch_descendant_clears_checkout_identity_and_resumes_as_legacy() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_credentials_effort_and_openrouter_key( + root.path(), + "gpt-5.4", + ProviderKind::OpenRouter, + crate::credentials::CredentialStorage::Memory, + None, + Some(crate::provider::OpenRouterApiKey::new("")), + ) + .unwrap(); + let parent = prepared(); + let mut claim = runtime.claim_session_fork().unwrap(); + let id = claim.id().to_owned(); + let error = runtime + .start_acp_driver_with_initial( + context(root.path()), + &mut claim, + Some(AcpForkState { + transcript: parent.clone(), + // Explicit ordinary-fork selection still wins over ancestry. + selection: ModelSelection::new(ProviderKind::OpenRouter, "invalid model"), + reasoning_effort: None, + parent_context: None, + }), + ) + .await + .err() + .expect("explicit invalid model must fail startup"); + assert!( + error + .to_string() + .contains("model name is outside canonical bounds"), + "{error}" + ); + let child = crate::session::load(root.path(), &id).unwrap(); + assert!(branch::BranchMetadata::read(&child).unwrap().is_none()); + assert!(branch::BranchMetadata::read(&parent).unwrap().is_some()); + // Retain this ordinary fork fixture to exercise its next load. It has + // normal creation authority, not a prompt-checkout completion record. + claim.commit().unwrap(); + let mut resume = runtime.claim_session_load(&id).unwrap(); + let error = runtime + .start_acp_driver(context(root.path()), &mut resume) + .await + .err() + .expect("legacy defaults reach the empty-key adapter error"); + assert!( + error + .to_string() + .contains("--openrouter-api-key cannot be empty"), + "{error}" + ); + } + + #[test] + fn ordinary_fork_rejects_malformed_ancestry_before_stripping() { + let mut source = prepared(); + source[0].metadata.insert("unrelated".into(), json!(true)); + let child = ordinary_acp_fork_transcript(source.clone()).unwrap(); + assert_eq!(child[0].metadata["unrelated"], json!(true)); + assert!(source[0].metadata.contains_key(branch::METADATA_KEY)); + source[0] + .metadata + .insert(branch::METADATA_KEY.into(), Value::Null); + assert!(ordinary_acp_fork_transcript(source).is_err()); + } + + #[test] + fn branch_activation_suppresses_only_the_committed_prompt_append() { + use agentkit_loop::{TranscriptEvent, TranscriptObserver}; + let root = tempfile::tempdir().unwrap(); + let opened = + crate::session::open_uncommitted(root.path(), "child", false, prepared()).unwrap(); + branch::commit(&opened.observer, &opened.transcript).unwrap(); + let prompt = opened.transcript.last().unwrap().clone(); + let observer = AcpTranscriptObserver { + observer: opened.observer, + committed_input: Mutex::new(Some(prompt.clone())), + }; + let session_id = agentkit_core::SessionId::new("child"); + observer.on_transcript_event(TranscriptEvent { + session_id: &session_id, + item: &prompt, + }); + assert_eq!( + crate::session::load(root.path(), "child").unwrap(), + opened.transcript + ); + // An identical later user prompt is a real append, never deduplicated. + observer.on_transcript_event(TranscriptEvent { + session_id: &session_id, + item: &prompt, + }); + let mut expected = opened.transcript; + expected.push(prompt); + assert_eq!( + crate::session::load(root.path(), "child").unwrap(), + expected + ); + } +} diff --git a/src/runtime/input_settlement.rs b/src/runtime/input_settlement.rs new file mode 100644 index 0000000..e739423 --- /dev/null +++ b/src/runtime/input_settlement.rs @@ -0,0 +1,758 @@ +//! Persist acknowledged steering without executing its cancelled continuation. + +use std::{ + ops::{Deref, DerefMut}, + sync::{Arc, Mutex}, +}; + +use agentkit_core::{FinishReason, Item}; +use agentkit_loop::{ + LoopCtx, LoopDriver, LoopError, LoopMutator, LoopStep, ModelSession, TranscriptCursor, + TranscriptEvent, TranscriptObserver, +}; +use async_trait::async_trait; + +/// A per-driver fence, registered FIRST, before compaction or any other mutator. +/// Clones share the fence between the registered mutator and its driver wrapper; +/// do not share it between drivers. +#[derive(Clone, Default)] +pub(crate) struct InputSettlement(Arc>); + +// Only the wrapper arms and RAII disarms. Idle has no exclusion; an armed +// fence is either full (no exclusion) or partial. The observer advances the +// pending position; the first mutator verifies and strips. No guard spans an +// await or an external callback. Poison isolates all users except reset-only Drop. +#[derive(Default)] +struct SettlementState { + active: bool, + exclusion: Option, +} + +/// Disposition of the raw queue, tracked by position rather than item identity. +struct ExcludedInput { + pending: Vec, + acknowledged: usize, + observed: usize, + transcript_len: usize, + rejected: bool, + stripped: bool, +} + +fn matches_stamped(raw: &Item, stamped: &Item) -> bool { + let mut normalized = stamped.clone(); + if raw.created_at.is_none() { + normalized.created_at = None; + } + raw == &normalized +} + +struct SettlementObserver { + settlement: InputSettlement, + inner: O, +} + +impl TranscriptObserver for SettlementObserver { + fn on_transcript_event(&self, event: TranscriptEvent<'_>) { + let forward = { + let Ok(mut state) = self.settlement.0.lock() else { + // An uncertain disposition must never reach persistence. + return; + }; + match state.exclusion.as_mut() { + None => true, + // The first mutator verified and removed the raw suffix. The + // pinned driver now appends its normal cancellation diagnostic. + Some(plan) if plan.stripped => true, + Some(plan) => { + if plan.rejected + || !plan + .pending + .get(plan.observed) + .is_some_and(|raw| matches_stamped(raw, event.item)) + { + plan.rejected = true; + false + } else { + let forward = plan.observed < plan.acknowledged; + plan.observed += 1; + forward + } + } + } + }; + // Persistence may panic or acquire locks. Never call it under the + // disposition lock; RAII isolates the driver on unwind. + if forward { + self.inner.on_transcript_event(event); + } + } +} + +impl InputSettlement { + /// Register this around the real persistence observer on the same driver. + pub(crate) fn observer( + &self, + inner: O, + ) -> impl TranscriptObserver + use { + SettlementObserver { + settlement: self.clone(), + inner, + } + } + + /// Wrap the driver on which a clone of this fence was registered first. + pub(crate) fn wrap(self, driver: LoopDriver) -> InputSettlingDriver { + InputSettlingDriver { + driver, + settlement: self, + settlement_failed: false, + } + } +} + +#[async_trait] +impl LoopMutator for InputSettlement { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + _ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + let mut state = self + .0 + .lock() + .map_err(|_| LoopError::InvalidState("input settlement state poisoned".into()))?; + if !state.active { + return Ok(()); + } + if let Some(plan) = state.exclusion.as_mut() { + if plan.rejected + || plan.stripped + || plan.observed != plan.pending.len() + || cursor.len() != plan.transcript_len + plan.pending.len() + || !plan + .pending + .iter() + .zip(&cursor[plan.transcript_len..]) + .all(|(raw, stamped)| matches_stamped(raw, stamped)) + { + return Err(LoopError::InvalidState( + "input settlement did not observe the expected pending tail".into(), + )); + } + cursor.truncate(plan.transcript_len + plan.acknowledged); + plan.stripped = true; + } + Err(LoopError::Cancelled) + } +} + +/// Owns the only arming scope. Drop resets all disposition state on errors, +/// unwind, and cancellation, without calling external code. +struct ArmedSettlement<'a>(&'a InputSettlement); + +impl Drop for ArmedSettlement<'_> { + fn drop(&mut self) { + let old = { + // Recovery is reset-only: discard EVERY field, never continue a + // possibly incomplete transition. The owning driver stays unavailable. + let mut state = self.0.0.lock().unwrap_or_else(|error| error.into_inner()); + std::mem::take(&mut *state) + }; + drop(old); + } +} + +pub(crate) struct InputSettlingDriver { + driver: LoopDriver, + settlement: InputSettlement, + // Set before settlement can suspend; only verified full success restores use. + settlement_failed: bool, +} + +impl InputSettlingDriver { + /// False after partial or incomplete settlement, including a dropped future. + /// Actor owners must retire unavailable drivers rather than drive them again. + pub(crate) fn is_available(&self) -> bool { + !self.settlement_failed + } + + /// Isolate a driver whose input acknowledgement boundary is uncertain. + pub(crate) fn make_unavailable(&mut self) { + self.settlement_failed = true; + } + + /// Transfer ownership to a protocol that does not settle injected input. + pub(crate) fn into_inner(self) -> LoopDriver { + self.driver + } +} + +impl Deref for InputSettlingDriver { + type Target = LoopDriver; + + fn deref(&self) -> &Self::Target { + &self.driver + } +} + +impl DerefMut for InputSettlingDriver { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.driver + } +} + +impl InputSettlingDriver { + /// Settle only successfully delivered steering after cancellation at an + /// injection boundary whose pending-input baseline was verified empty. + /// Original unstarted prompts and synthetic compaction input must not use + /// this path. The caller must stop driving on error or an aborted future. + /// + /// Retirement removes foreground resumptions before the single fenced + /// step. That step appends input through the normal transcript observer, + /// then the first mutator cancels before compaction or model execution. + /// It produces a fresh cancelled logical turn and its normal diagnostic. + pub(crate) async fn settle_delivered_input(&mut self) -> Result<(), LoopError> { + self.settle_input(None).await + } + + /// Settle an exactly acknowledged prefix of the raw pending queue. The caller + /// supplies the original submitted items in order (including duplicates). + /// Requires the same empty injection-boundary baseline as full settlement. + /// An empty prefix discards the entire queue. Register `settlement.observer` + /// around persistence. A proper partial settlement leaves this driver + /// unavailable even on success: the actor MUST retire it, never retry the suffix. + pub(crate) async fn settle_acknowledged_input( + &mut self, + acknowledged: &[Item], + ) -> Result<(), LoopError> { + self.settle_input(Some(acknowledged)).await + } + + async fn settle_input(&mut self, acknowledged: Option<&[Item]>) -> Result<(), LoopError> { + if !self.is_available() { + return Err(LoopError::InvalidState( + "driver is unavailable after partial or incomplete input settlement".into(), + )); + } + self.settlement_failed = true; + let pending = self.driver.snapshot().pending_input; + if pending.is_empty() { + return Err(LoopError::InvalidState( + "input settlement requires delivered pending input".into(), + )); + } + let acknowledged = acknowledged.unwrap_or(&pending); + if !pending.starts_with(acknowledged) { + return Err(LoopError::InvalidState( + "acknowledged input is not an exact pending prefix".into(), + )); + } + let partial = acknowledged.len() < pending.len(); + let acknowledged_count = acknowledged.len(); + self.driver.retire_interrupted_turn().await?; + let exclusion = partial.then(|| ExcludedInput { + pending, + acknowledged: acknowledged_count, + observed: 0, + transcript_len: self.driver.snapshot().transcript.len(), + rejected: false, + stripped: false, + }); + { + let mut state = + self.settlement.0.lock().map_err(|_| { + LoopError::InvalidState("input settlement state poisoned".into()) + })?; + if state.active { + return Err(LoopError::InvalidState( + "input settlement fence is already armed".into(), + )); + } + *state = SettlementState { + active: true, + exclusion, + }; + } + let _armed = ArmedSettlement(&self.settlement); + let step = self.driver.next().await?; + if !matches!(step, LoopStep::Finished(turn) if turn.finish_reason == FinishReason::Cancelled) + || !self.driver.snapshot().pending_input.is_empty() + { + return Err(LoopError::InvalidState( + "input settlement did not finish cancelled with empty pending input".into(), + )); + } + if partial { + let state = + self.settlement.0.lock().map_err(|_| { + LoopError::InvalidState("input settlement state poisoned".into()) + })?; + if !state + .exclusion + .as_ref() + .is_some_and(|plan| plan.stripped && !plan.rejected) + { + return Err(LoopError::InvalidState( + "input settlement exclusion was not verified".into(), + )); + } + } + self.settlement_failed = partial; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use agentkit_core::{Item, ItemKind, TurnCancellation}; + use agentkit_loop::{ + Agent, ModelAdapter, ModelTurn, ModelTurnEvent, SessionConfig, TurnRequest, + }; + + use super::*; + + struct UnavailableProvider; + + #[async_trait] + impl ModelAdapter for UnavailableProvider { + type Session = Self; + + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(Self) + } + } + + #[async_trait] + impl ModelSession for UnavailableProvider { + type Turn = Self; + + async fn begin_turn( + &mut self, + _: TurnRequest, + _: Option, + ) -> Result { + Err(LoopError::Provider("provider unavailable".into())) + } + } + + #[async_trait] + impl ModelTurn for UnavailableProvider { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + unreachable!("the unavailable provider cannot start a turn") + } + } + + struct UnavailableCompactor; + + #[async_trait] + impl LoopMutator for UnavailableCompactor { + async fn mutate( + &self, + _: &mut TranscriptCursor<'_>, + _: LoopCtx<'_>, + ) -> Result<(), LoopError> { + Err(LoopError::Mutator("compactor unavailable".into())) + } + } + + /// An unavailable task service can either fail its update read or leave it + /// pending. All other task operations use the real synchronous manager. + struct UnavailableTaskService { + inner: agentkit_task_manager::SimpleTaskManager, + pending: bool, + } + + #[async_trait] + impl agentkit_task_manager::TaskManager for UnavailableTaskService { + async fn start_task( + &self, + request: agentkit_task_manager::TaskLaunchRequest, + ctx: agentkit_task_manager::TaskStartContext, + ) -> Result + { + self.inner.start_task(request, ctx).await + } + + async fn wait_for_turn( + &self, + turn_id: &agentkit_core::TurnId, + cancellation: Option, + ) -> Result< + Option, + agentkit_task_manager::TaskManagerError, + > { + self.inner.wait_for_turn(turn_id, cancellation).await + } + + async fn take_pending_loop_updates( + &self, + ) -> Result< + agentkit_task_manager::PendingLoopUpdates, + agentkit_task_manager::TaskManagerError, + > { + if self.pending { + std::future::pending().await + } else { + Err(agentkit_task_manager::TaskManagerError::Internal( + "task service unavailable".into(), + )) + } + } + + async fn on_turn_interrupted( + &self, + turn_id: &agentkit_core::TurnId, + ) -> Result<(), agentkit_task_manager::TaskManagerError> { + self.inner.on_turn_interrupted(turn_id).await + } + + fn handle(&self) -> agentkit_task_manager::TaskManagerHandle { + self.inner.handle() + } + } + + #[tokio::test] + async fn failed_or_aborted_settlement_keeps_driver_unavailable() { + use std::{future::Future, task::Poll}; + + for pending in [false, true] { + let settlement = InputSettlement::default(); + let raw = Agent::builder() + .model(UnavailableProvider) + .mutator(settlement.clone()) + .task_manager(UnavailableTaskService { + inner: agentkit_task_manager::SimpleTaskManager::new(), + pending, + }) + .build() + .unwrap() + .start(SessionConfig::new("failed-settlement")) + .await + .unwrap(); + let mut driver = settlement.clone().wrap(raw); + driver + .submit_input(vec![Item::text(ItemKind::User, "delivered steering")]) + .unwrap(); + assert!(driver.is_available()); + if pending { + let mut settling = Box::pin(driver.settle_delivered_input()); + std::future::poll_fn(|cx| { + assert!(settling.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(settling); + } else { + assert!(matches!( + driver.settle_delivered_input().await, + Err(LoopError::Tool(_)) + )); + } + assert!(!driver.is_available()); + assert!(!settlement.0.lock().unwrap().active); + assert!(matches!( + driver.settle_delivered_input().await, + Err(LoopError::InvalidState(_)) + )); + assert!(!driver.is_available()); + } + } + + fn persisted_records(storage: &std::path::Path, id: &str) -> String { + let directory = std::fs::read_dir(storage) + .unwrap() + .map(|entry| entry.unwrap().path()) + .find(|path| path.is_dir()) + .unwrap(); + std::fs::read_to_string(directory.join(format!("{id}.jsonl"))).unwrap() + } + + #[tokio::test] + async fn acknowledged_prefix_persists_by_position_without_temporary_suffix() { + for stamped in [false, true] { + for accepted in [0, 1, 2, 3, 4] { + let root = tempfile::tempdir().unwrap(); + let storage = tempfile::tempdir().unwrap(); + let id = "partial-settlement"; + let opened = crate::session::open_in( + root.path(), + storage.path(), + id, + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let settlement = InputSettlement::default(); + let raw = Agent::builder() + .model(UnavailableProvider) + .mutator(settlement.clone()) + .mutator(UnavailableCompactor) + .transcript(opened.transcript) + .transcript_observer(settlement.observer(opened.observer)) + .build() + .unwrap() + .start(SessionConfig::new(id)) + .await + .unwrap(); + let mut driver = settlement.clone().wrap(raw); + let mut same = Item::text(ItemKind::User, "queued-identical"); + if stamped { + same.created_at = Some(agentkit_core::Timestamp::now()); + } + let pending = vec![ + same.clone(), + same.clone(), + same, + Item::text(ItemKind::User, "unacknowledged-distinct"), + ]; + driver.submit_input(pending.clone()).unwrap(); + driver + .settle_acknowledged_input(&pending[..accepted]) + .await + .unwrap(); + assert_eq!(driver.is_available(), accepted == pending.len()); + let snapshot = driver.snapshot(); + assert!(snapshot.pending_input.is_empty()); + assert_eq!(snapshot.transcript.len(), 2 + accepted); + for (raw, actual) in pending[..accepted] + .iter() + .zip(&snapshot.transcript[1..1 + accepted]) + { + assert!(matches_stamped(raw, actual)); + assert!(actual.created_at.is_some()); + } + assert_eq!( + crate::session::load_in(root.path(), storage.path(), id).unwrap(), + snapshot.transcript + ); + // Inspect append history, not only the canonical reload: a later + // replacement must not conceal a temporarily persisted suffix. + let records = persisted_records(storage.path(), id); + assert_eq!(records.matches("queued-identical").count(), accepted.min(3)); + assert_eq!(records.contains("unacknowledged-distinct"), accepted == 4); + { + let state = settlement.0.lock().unwrap(); + assert!(!state.active && state.exclusion.is_none()); + } + if accepted < pending.len() { + assert!(driver.settle_delivered_input().await.is_err()); + assert_eq!(persisted_records(storage.path(), id), records); + } + } + } + } + + #[tokio::test] + async fn partial_failure_abort_and_wrong_prefix_never_persist_suffix() { + use std::{future::Future, task::Poll}; + + for mode in [ + "failure", + "abort", + "wrong-prefix", + "too-long", + "missing-observer", + ] { + let root = tempfile::tempdir().unwrap(); + let storage = tempfile::tempdir().unwrap(); + let id = "failed-partial"; + let opened = crate::session::open_in( + root.path(), + storage.path(), + id, + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let settlement = InputSettlement::default(); + let builder = Agent::builder() + .model(UnavailableProvider) + .mutator(settlement.clone()) + .transcript(opened.transcript); + // Missing filtering registration must fail closed at the mutator. + let builder = if mode == "missing-observer" { + builder + } else { + builder + .transcript_observer(settlement.observer(opened.observer.clone())) + .task_manager(UnavailableTaskService { + inner: agentkit_task_manager::SimpleTaskManager::new(), + pending: mode == "abort", + }) + }; + let raw = builder + .build() + .unwrap() + .start(SessionConfig::new(id)) + .await + .unwrap(); + let mut driver = settlement.clone().wrap(raw); + let a = Item::text(ItemKind::User, "acknowledged"); + let b = Item::text(ItemKind::User, "excluded"); + let pending = vec![a.clone(), b.clone()]; + driver.submit_input(pending.clone()).unwrap(); + let witness = match mode { + "wrong-prefix" => vec![b], + "too-long" => vec![a.clone(), b, a], + _ => vec![a], + }; + if mode == "abort" { + let mut settling = Box::pin(driver.settle_acknowledged_input(&witness)); + std::future::poll_fn(|cx| { + assert!(settling.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(settling); + } else { + assert!(driver.settle_acknowledged_input(&witness).await.is_err()); + } + assert!(!driver.is_available()); + { + let state = settlement.0.lock().unwrap(); + assert!(!state.active && state.exclusion.is_none()); + } + let records = persisted_records(storage.path(), id); + assert!(!records.contains("acknowledged") && !records.contains("excluded")); + assert_eq!( + crate::session::load_in(root.path(), storage.path(), id) + .unwrap() + .len(), + 1 + ); + if mode != "missing-observer" { + assert_eq!(driver.snapshot().pending_input, pending); + } + assert!(driver.settle_acknowledged_input(&witness).await.is_err()); + assert_eq!(persisted_records(storage.path(), id), records); + } + } + + struct PanicAfterPersist(O); + + impl TranscriptObserver for PanicAfterPersist { + fn on_transcript_event(&self, event: TranscriptEvent<'_>) { + self.0.on_transcript_event(event); + panic!("observer failed after persistence"); + } + } + + #[tokio::test] + async fn partial_observer_unwind_disarms_without_poison_or_retry() { + use futures_util::FutureExt; + + let root = tempfile::tempdir().unwrap(); + let storage = tempfile::tempdir().unwrap(); + let id = "unwind-partial"; + let opened = crate::session::open_in( + root.path(), + storage.path(), + id, + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let settlement = InputSettlement::default(); + let raw = Agent::builder() + .model(UnavailableProvider) + .mutator(settlement.clone()) + .transcript(opened.transcript) + .transcript_observer(settlement.observer(PanicAfterPersist(opened.observer))) + .build() + .unwrap() + .start(SessionConfig::new(id)) + .await + .unwrap(); + let mut driver = settlement.clone().wrap(raw); + let a = Item::text(ItemKind::User, "acknowledged"); + driver + .submit_input(vec![a.clone(), Item::text(ItemKind::User, "excluded")]) + .unwrap(); + assert!( + std::panic::AssertUnwindSafe(driver.settle_acknowledged_input(&[a])) + .catch_unwind() + .await + .is_err() + ); + assert!(!driver.is_available()); + { + let state = settlement.0.lock().unwrap(); + assert!(!state.active && state.exclusion.is_none()); + } + let records = persisted_records(storage.path(), id); + assert_eq!(records.matches("acknowledged").count(), 1); + assert!(!records.contains("excluded")); + assert_eq!( + crate::session::load_in(root.path(), storage.path(), id) + .unwrap() + .len(), + 2 + ); + assert!(driver.settle_delivered_input().await.is_err()); + assert_eq!(persisted_records(storage.path(), id), records); + } + + #[tokio::test] + async fn settlement_persists_without_compactor_or_provider_and_disarms() { + for with_compactor in [false, true] { + let root = tempfile::tempdir().unwrap(); + let storage = tempfile::tempdir().unwrap(); + let id = "input-settlement"; + let opened = crate::session::open_in( + root.path(), + storage.path(), + id, + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let settlement = InputSettlement::default(); + let builder = Agent::builder() + .model(UnavailableProvider) + .mutator(settlement.clone()) + .transcript(opened.transcript) + .transcript_observer(settlement.observer(opened.observer)); + let builder = if with_compactor { + builder.mutator(UnavailableCompactor) + } else { + builder + }; + let raw = builder + .build() + .unwrap() + .start(SessionConfig::new(id)) + .await + .unwrap(); + let mut driver = settlement.wrap(raw); + let delivered = Item::text(ItemKind::User, "delivered steering"); + driver.submit_input(vec![delivered.clone()]).unwrap(); + driver.settle_delivered_input().await.unwrap(); + assert!(driver.is_available()); + let snapshot = driver.snapshot(); + assert!(snapshot.pending_input.is_empty()); + assert_eq!(snapshot.transcript[1].parts, delivered.parts); + assert!(snapshot.transcript[1].created_at.is_some()); + assert_eq!( + crate::session::load_in(root.path(), storage.path(), id).unwrap(), + snapshot.transcript + ); + + driver + .submit_input(vec![Item::text(ItemKind::User, "fresh prompt")]) + .unwrap(); + let error = driver.next().await.unwrap_err(); + if with_compactor { + assert!(matches!(error, LoopError::Mutator(_))); + } else { + assert!(matches!(error, LoopError::Provider(_))); + } + } + } +} diff --git a/src/session.rs b/src/session.rs index a749a7c..62c58d5 100644 --- a/src/session.rs +++ b/src/session.rs @@ -21,6 +21,8 @@ use agentkit_core::{Item, ItemKind, Part, Timestamp}; use agentkit_loop::{TranscriptEvent, TranscriptObserver}; use serde::{Deserialize, Serialize}; +pub(crate) mod branch; + pub const SCHEMA_VERSION: u32 = 3; const REDIRECT_SCHEMA_VERSION: u32 = 4; const PREVIOUS_SCHEMA_VERSION: u32 = 2; @@ -211,6 +213,7 @@ pub(crate) fn clone_completed_in( destination: &str, ) -> Result<(), String> { let mut transcript = load_in(root, directory, source)?; + branch::strip_for_plain_fork(&mut transcript)?; crate::transcript::sanitize_forked_transcript(&mut transcript); let opened = open_with_initial_timestamps_in( root, @@ -316,6 +319,9 @@ pub(crate) fn open_uncommitted( force: bool, initial: Vec, ) -> Result { + // Checkout fingerprints bind exact historical items, including unknown + // timestamps. The new submitted prompt is stamped during preparation. + let stamp_items = branch::BranchMetadata::read(&initial)?.is_none(); open_with_initial_timestamps_in( root, &default_directory()?, @@ -324,7 +330,7 @@ pub(crate) fn open_uncommitted( force, initial, InitialTranscriptOptions { - stamp_items: true, + stamp_items, commit_creation: false, }, ) @@ -757,6 +763,9 @@ struct TranscriptHistory { items: Vec, generation: u64, states: Vec>, + // Each replacement is a prefix of a retained state. Keep its exact length + // without duplicating every historical transcript in memory. + replacement_boundaries: Vec<(usize, usize)>, } enum StoredTranscript { @@ -818,6 +827,7 @@ fn read_record_lines( let mut items = Vec::new(); let mut expected = 1_u64; let mut states = Vec::new(); + let mut replacement_boundaries = Vec::new(); let mut redirect = None; for (index, line) in lines.enumerate() { let line = line?; @@ -860,6 +870,7 @@ fn read_record_lines( if !items.is_empty() { states.push(items.clone()); } + replacement_boundaries.push((states.len(), replacement.len())); items = replacement; } (None, None, Some(target)) @@ -888,6 +899,7 @@ fn read_record_lines( items, generation: expected - 1, states, + replacement_boundaries, })) } @@ -1339,7 +1351,10 @@ fn belongs_to_workspace_in( ) -> Result { validate_id(session_id)?; let root = canonical_workspace(root); - Ok(select_authority(global_directory, &root, session_id)?.is_some()) + // ACP resume checks workspace membership before acquiring the writer. + // Apply the same read-only recovery view as resume preflight; the locked + // opener alone may repair a torn migration destination or tail. + Ok(select_authority_with(global_directory, &root, session_id, true, true)?.is_some()) } pub(crate) fn list_ids_in(directory: &Path) -> Result, String> { @@ -1408,20 +1423,8 @@ fn select_authority( fn transcript_snapshot(path: &Path, tolerate_incomplete_tail: bool) -> Result, String> { let mut bytes = fs::read(path).map_err(|error| format!("could not read {}: {error}", path.display()))?; - if tolerate_incomplete_tail && !bytes.ends_with(b"\n") { - let tail_start = bytes - .iter() - .rposition(|byte| *byte == b'\n') - .map_or(0, |index| index + 1); - let tail = &bytes[tail_start..]; - let incomplete = match std::str::from_utf8(tail) { - Ok(_) => serde_json::from_slice::(tail) - .is_err_and(|error| error.classify() == serde_json::error::Category::Eof), - Err(error) => error.error_len().is_none(), - }; - if incomplete { - bytes.truncate(tail_start); - } + if tolerate_incomplete_tail && let Some(complete) = torn_migration_tail_start(&bytes) { + bytes.truncate(complete); } Ok(bytes) } @@ -1467,6 +1470,16 @@ fn select_authority_with( { continue; } + // The locked opener discards an empty or torn first replacement at + // the scoped destination. Ignore exactly that destination here too; + // legacy candidates still undergo normal workspace/lineage validation. + if tolerate_incomplete_tail && path == &scoped { + let bytes = fs::read(path) + .map_err(|error| format!("could not read {}: {error}", path.display()))?; + if incomplete_migration_destination(&bytes) { + continue; + } + } let (workspace, transcript) = read_authority_candidate(path, session_id, tolerate_incomplete_tail)?; if is_global @@ -1574,11 +1587,23 @@ fn torn_migration_tail_start(bytes: &[u8]) -> Option { .iter() .rposition(|byte| *byte == b'\n') .map_or(0, |index| index + 1); - serde_json::from_slice::(&bytes[start..]) - .is_err() + // Missing fields, invalid syntax, and unknown record shapes are not torn + // writes. Even without a newline, a complete JSON value must be validated. + let tail = &bytes[start..]; + if std::str::from_utf8(tail).is_err_and(|error| error.error_len().is_some()) { + return None; + } + serde_json::from_slice::(tail) + .is_err_and(|error| error.is_eof()) .then_some(start) } +/// No complete record exists yet; only a scoped migration destination may be +/// discarded. Complete malformed records must still reach the strict parser. +fn incomplete_migration_destination(bytes: &[u8]) -> bool { + bytes.is_empty() || torn_migration_tail_start(bytes) == Some(0) +} + fn migration_source_workspace(path: &Path, session_id: &str) -> Result, String> { let bytes = fs::read(path).map_err(|error| format!("could not read {}: {error}", path.display()))?; @@ -1629,7 +1654,7 @@ fn recover_torn_migration_writes( } let bytes = fs::read(&path) .map_err(|error| format!("could not read {}: {error}", path.display()))?; - if bytes.is_empty() && path == scoped { + if path == scoped && incomplete_migration_destination(&bytes) { filesystem .remove_file(&path) .map_err(|error| format!("could not remove {}: {error}", path.display()))?; @@ -1637,27 +1662,47 @@ fn recover_torn_migration_writes( continue; } let Some(complete) = torn_migration_tail_start(&bytes) else { + if !bytes.is_empty() && !bytes.ends_with(b"\n") { + // A complete record can lose only its separator. Validate the + // entire history before repairing it under the mutation lock, + // before either migration or the resumed writer can append. + read_records_bytes(&path, session_id, &bytes)?; + if let Some(workspace) = transcript_workspace_bytes(&path, session_id, &bytes)? + && workspace != root + { + return Err(format!( + "session {session_id:?} belongs to workspace {}, not {}", + workspace.display(), + root.display() + )); + } + let mut file = OpenOptions::new() + .append(true) + .open_in(&filesystem, &path) + .map_err(|error| format!("could not open {}: {error}", path.display()))?; + file.write_all(b"\n") + .and_then(|_| file.sync_all()) + .map_err(|error| { + format!( + "could not recover transcript separator {}: {error}", + path.display() + ) + })?; + } continue; }; - if complete == 0 && path == scoped { - filesystem - .remove_file(&path) - .map_err(|error| format!("could not remove {}: {error}", path.display()))?; - sync_parent_directory(&filesystem, &path)?; - } else { - let file = OpenOptions::new() - .write(true) - .open_in(&filesystem, &path) - .map_err(|error| format!("could not open {}: {error}", path.display()))?; - file.set_len(complete as u64) - .and_then(|_| file.sync_all()) - .map_err(|error| { - format!( - "could not recover torn migration {}: {error}", - path.display() - ) - })?; - } + let file = OpenOptions::new() + .write(true) + .open_in(&filesystem, &path) + .map_err(|error| format!("could not open {}: {error}", path.display()))?; + file.set_len(complete as u64) + .and_then(|_| file.sync_all()) + .map_err(|error| { + format!( + "could not recover torn migration {}: {error}", + path.display() + ) + })?; } Ok(()) } @@ -1936,6 +1981,14 @@ fn workspace_storage_directory(directory: &Path, root: &Path) -> PathBuf { directory.join(format!("w-{}", identity.to_hex())) } +#[cfg(test)] +pub(crate) fn transcript_path_for_test(root: &Path, session_id: &str) -> PathBuf { + transcript_path( + &workspace_storage_directory(&default_directory().unwrap(), &canonical_workspace(root)), + session_id, + ) +} + fn transcript_path(directory: &Path, session_id: &str) -> PathBuf { directory.join(format!("{session_id}.jsonl")) } @@ -1965,6 +2018,38 @@ pub(crate) fn validate_id(value: &str) -> Result<(), String> { } } +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + + /// Leave an intact legacy source and a crashed first scoped replacement. + pub(crate) fn interrupted_migration( + root: &Path, + session_id: &str, + transcript: Vec, + empty: bool, + ) -> (PathBuf, PathBuf) { + let legacy = legacy_transcript(root, session_id); + let scoped = transcript_path_for_test(root, session_id); + let record = Record { + schema_version: SCHEMA_VERSION, + session_id: session_id.into(), + generation: 1, + workspace_root: Some(canonical_workspace(root)), + item: None, + replacement: Some(transcript), + redirect: None, + }; + let mut bytes = serde_json::to_vec(&record).unwrap(); + bytes.push(b'\n'); + std::fs::create_dir_all(legacy.parent().unwrap()).unwrap(); + std::fs::create_dir_all(scoped.parent().unwrap()).unwrap(); + std::fs::write(&legacy, &bytes).unwrap(); + std::fs::write(&scoped, &bytes[..if empty { 0 } else { bytes.len() / 2 }]).unwrap(); + (legacy, scoped) + } +} + #[cfg(test)] mod tests { use std::io::{BufRead, BufReader}; @@ -2800,6 +2885,113 @@ mod tests { )); } + #[test] + fn torn_tail_requires_json_eof_and_only_incomplete_final_utf8() { + for tail in [b"{".as_slice(), b"{\"item\":", b"{\"item\":\"\xe2\x82"] { + let mut bytes = b"complete\n".to_vec(); + bytes.extend_from_slice(tail); + assert_eq!(torn_migration_tail_start(&bytes), Some(9), "{tail:?}"); + } + for tail in [ + b"{}".as_slice(), + b"{bad", + b"{bad\xe2", + b"{\"item\":\xe2", // a UTF-8 value cannot start outside a JSON string + b"{\"item\":\"\xff", // definite invalid UTF-8, not an interrupted codepoint + b"{\"item\":\"\xe2\x82\n", // incomplete nonfinal record + ] { + let mut bytes = b"complete\n".to_vec(); + bytes.extend_from_slice(tail); + assert_eq!(torn_migration_tail_start(&bytes), None, "{tail:?}"); + } + } + + #[test] + fn resume_repairs_only_missing_separator_before_appending() { + for remove_newline in [false, true] { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let mut expected = opened.transcript.clone(); + drop(opened); + let path = transcript_path(root.path(), "abc"); + let complete = fs::read(&path).unwrap(); + let input = if remove_newline { + &complete[..complete.len() - 1] + } else { + &complete[..] + }; + fs::write(&path, input).unwrap(); + assert_eq!(load(root.path(), "abc").unwrap(), expected); + assert_eq!(fs::read(&path).unwrap(), input); + let resumed = open(root.path(), "abc", true, false, Vec::new()).unwrap(); + assert_eq!(resumed.transcript, expected); + assert_eq!(fs::read(&path).unwrap(), complete); + let appended = + Item::text(ItemKind::User, "after resume").with_created_at(Timestamp(123)); + resumed.observer.on_transcript_event(TranscriptEvent { + session_id: &agentkit_core::SessionId::new("abc"), + item: &appended, + }); + expected.push(appended); + drop(resumed); + assert_eq!(load(root.path(), "abc").unwrap(), expected); + assert_eq!(fs::read_to_string(&path).unwrap().lines().count(), 2); + } + } + + #[test] + fn missing_separator_does_not_allow_invalid_history_to_be_repaired() { + for field in ["generation", "schema_version", "workspace_root"] { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + drop(opened); + let path = transcript_path(root.path(), "abc"); + let mut record: serde_json::Value = + serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + record[field] = if field == "workspace_root" { + serde_json::json!(root.path().join("other-workspace")) + } else { + serde_json::json!(999) + }; + let bytes = serde_json::to_vec(&record).unwrap(); + fs::write(&path, &bytes).unwrap(); + assert!(open(root.path(), "abc", true, false, Vec::new()).is_err()); + assert_eq!(fs::read(&path).unwrap(), bytes, "{field}"); + } + } + + #[test] + fn missing_legacy_separator_is_repaired_before_redirect_append() { + let root = tempfile::tempdir().unwrap(); + let global = super::transcript_path(&session_directory(root.path()), "abc"); + let expected = write_history(&global, PREVIOUS_SCHEMA_VERSION, "abc", &["legacy"], None); + let complete = fs::read(&global).unwrap(); + fs::write(&global, &complete[..complete.len() - 1]).unwrap(); + let resumed = open(root.path(), "abc", true, false, Vec::new()).unwrap(); + assert_eq!(resumed.transcript, expected); + drop(resumed); + assert_eq!(load(root.path(), "abc").unwrap(), expected); + assert!(fs::read(&global).unwrap().starts_with(&complete)); + assert!(matches!( + read_records_direct(&global, "abc").unwrap(), + StoredTranscript::Redirect(_) + )); + } + #[test] fn torn_new_scoped_authority_is_removed_and_recreated_from_legacy() { let root = tempfile::tempdir().unwrap(); @@ -2817,7 +3009,25 @@ mod tests { redirect: None, }; let encoded = serde_json::to_vec(&record).unwrap(); - fs::write(&scoped, &encoded[..encoded.len() / 2]).unwrap(); + let legacy_bytes = fs::read(&global).unwrap(); + for destination in [&b""[..], &encoded[..encoded.len() / 2]] { + fs::write(&scoped, destination).unwrap(); + let authority = select_authority_with( + &session_directory(root.path()), + &canonical_workspace(&project_root(root.path())), + "abc", + true, + true, + ) + .unwrap() + .unwrap(); + assert_eq!(authority.items, expected); + assert_eq!(authority.path, global); + assert_eq!(fs::read(&scoped).unwrap(), destination); + assert_eq!(fs::read(&global).unwrap(), legacy_bytes); + assert!(!scoped.with_extension("lock").exists()); + assert!(!global.with_extension("lock").exists()); + } assert!(load(root.path(), "abc").is_err()); let opened = open(root.path(), "abc", true, false, Vec::new()).unwrap(); diff --git a/src/session/branch.rs b/src/session/branch.rs new file mode 100644 index 0000000..0226bd7 --- /dev/null +++ b/src/session/branch.rs @@ -0,0 +1,1213 @@ +//! Versioned checkout provenance inside the existing bootstrap metadata. +//! +//! Metadata describes an intended branch, not a successful creation. Only a +//! complete, matching replacement record is completion evidence. Reads use the +//! normal authority selection/record validation and never rewrite history. + +use super::*; +use crate::{ReasoningEffort, provider::ModelSelection}; + +pub(crate) const METADATA_KEY: &str = "dev.kit.session.prompt_checkout"; +const VERSION: u32 = 1; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct Boundary { + /// Index in `load_history`, including states preceding compaction. + pub state_index: usize, + /// Number of items retained before the selected prompt. + pub prefix_len: usize, + /// BLAKE3 of the original parent prefix, including its metadata/timestamps. + pub prefix_hash: String, +} + +impl Boundary { + pub(crate) fn new(state_index: usize, prefix: &[Item]) -> Result { + if prefix.is_empty() { + return Err("checkout boundary requires a bootstrap prefix".into()); + } + Ok(Self { + state_index, + prefix_len: prefix.len(), + prefix_hash: digest(prefix)?, + }) + } +} + +/// Canonical IDs from the actual runtime types, not UI labels or adapter defaults. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct CapturedSelection { + pub provider: String, + pub model: String, + /// `default` is explicit: absence is not silently interpreted as default. + pub reasoning: String, +} + +impl CapturedSelection { + pub(crate) fn new(model: &ModelSelection, reasoning: Option) -> Self { + Self { + provider: model.provider.as_str().into(), + model: model.model.clone(), + reasoning: reasoning.map_or("default", ReasoningEffort::as_str).into(), + } + } + + pub(crate) fn resolve(&self) -> Result<(ModelSelection, Option), String> { + Ok(( + ModelSelection::from_id(&format!("{}:{}", self.provider, self.model))?, + ReasoningEffort::from_id(&self.reasoning)?, + )) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct SubmittedRequest { + pub id: String, + pub selection: CapturedSelection, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct Completion { + pub prefix_len: usize, + /// Hash of every retained item, with only this payload removed. + pub prefix_hash: String, + /// Hash of the full submitted prompt Item, including attachments/metadata. + pub prompt_hash: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct BranchMetadata { + pub version: u32, + pub parent_session_id: String, + pub boundary: Boundary, + pub checkout_id: String, + pub request: SubmittedRequest, + pub completion: Completion, +} + +impl BranchMetadata { + fn validate(&self) -> Result<(), String> { + if self.version != VERSION { + return Err(format!( + "unsupported prompt checkout metadata version {}", + self.version + )); + } + validate_id(&self.parent_session_id)?; + if self.checkout_id.trim().is_empty() || self.request.id.trim().is_empty() { + return Err( + "prompt checkout and submitted request identities must not be empty".into(), + ); + } + if self.boundary.prefix_len == 0 || self.completion.prefix_len != self.boundary.prefix_len { + return Err("invalid prompt checkout prefix length".into()); + } + for hash in [ + &self.boundary.prefix_hash, + &self.completion.prefix_hash, + &self.completion.prompt_hash, + ] { + if hash.len() != 64 + || !hash + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + return Err("invalid prompt checkout content hash".into()); + } + } + self.request.selection.resolve()?; + Ok(()) + } + + /// Missing metadata is valid for legacy/root sessions. Present invalid data + /// is always an error, including explicit null and unknown nested fields. + pub(crate) fn read(transcript: &[Item]) -> Result, String> { + if transcript + .iter() + .skip(1) + .any(|item| item.metadata.contains_key(METADATA_KEY)) + { + return Err("prompt checkout metadata must be on the bootstrap item".into()); + } + let Some(value) = transcript + .first() + .and_then(|item| item.metadata.get(METADATA_KEY)) + else { + return Ok(None); + }; + if transcript[0].kind != ItemKind::System { + return Err("prompt checkout metadata requires a system bootstrap".into()); + } + let version = value + .get("version") + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| "invalid prompt checkout metadata version".to_string())?; + if version != u64::from(VERSION) { + return Err(format!( + "unsupported prompt checkout metadata version {version}" + )); + } + let metadata: Self = serde_json::from_value(value.clone()) + .map_err(|error| format!("invalid prompt checkout metadata: {error}"))?; + metadata.validate()?; + Ok(Some(metadata)) + } + + fn matches_snapshot(&self, transcript: &[Item]) -> Result { + if transcript.len().checked_sub(1) != Some(self.completion.prefix_len) { + return Ok(false); + } + let (prompt, prefix) = transcript.split_last().expect("nonempty checked snapshot"); + Ok(prompt.kind == ItemKind::User + && digest(prompt)? == self.completion.prompt_hash + && prefix_digest(prefix)? == self.completion.prefix_hash) + } +} + +/// Validate a branch reload before an opener can repair history or an adapter +/// can start. Root/legacy transcripts retain their ordinary defaults. +pub(crate) fn validate_resume( + root: &Path, + session_id: &str, +) -> Result, String> { + validate_id(session_id)?; + // Root/legacy openers already recover torn migration tails while holding + // their locks. Inspect a bounded complete prefix here without rewriting it, + // so a missing checkout payload does not disable that existing recovery. + let authority = select_authority_with( + &default_directory()?, + &canonical_workspace(root), + session_id, + true, + true, + )? + .ok_or_else(|| format!("session {session_id:?} does not exist"))?; + for state in &authority.historical_items { + BranchMetadata::read(state)?; + } + let Some(metadata) = BranchMetadata::read(&authority.items)? else { + return Ok(None); + }; + let committed = lookup_committed( + root, + session_id, + &metadata.checkout_id, + &metadata.request.id, + )? + .ok_or_else(|| format!("session {session_id:?} has an incomplete prompt checkout"))?; + if committed.metadata != metadata { + return Err(format!( + "session {session_id:?} prompt checkout metadata differs from its completion" + )); + } + Ok(Some(metadata)) +} + +/// A plain fork has no checkout completion of its own. Validate inherited +/// metadata before removing its request identity from the new child only. +pub(crate) fn strip_for_plain_fork(transcript: &mut [Item]) -> Result<(), String> { + BranchMetadata::read(transcript)?; + if let Some(bootstrap) = transcript.first_mut() { + bootstrap.metadata.remove(METADATA_KEY); + } + Ok(()) +} + +/// Prepare the exact initial transcript to pass to `open_uncommitted`. +/// The caller selects a validated boundary and sanitizes session-bound provider +/// continuation metadata before this call. The inherited checkout payload is +/// replaced, never merged. Historical unknown timestamps remain unknown. +pub(crate) fn prepare( + mut prefix: Vec, + parent_session_id: String, + boundary: Boundary, + checkout_id: String, + request: SubmittedRequest, + mut prompt: Item, +) -> Result, String> { + if prefix.is_empty() || prefix[0].kind != ItemKind::System || prompt.kind != ItemKind::User { + return Err("prompt checkout requires a system bootstrap and a user prompt".into()); + } + // Validate inherited data before replacing it; malformed ancestry is not legacy. + BranchMetadata::read(&prefix)?; + stamp_item(&mut prompt, Timestamp::now()); + let metadata = BranchMetadata { + version: VERSION, + parent_session_id, + boundary, + checkout_id, + request, + completion: Completion { + prefix_len: prefix.len(), + prefix_hash: prefix_digest(&prefix)?, + prompt_hash: digest(&prompt)?, + }, + }; + metadata.validate()?; + prefix[0].metadata.insert( + METADATA_KEY.into(), + serde_json::to_value(metadata) + .map_err(|error| format!("could not encode prompt checkout metadata: {error}"))?, + ); + prefix.push(prompt); + Ok(prefix) +} + +fn digest(value: &(impl Serialize + ?Sized)) -> Result { + // Going through Value canonicalizes map order before hashing, independent of + // runtime MetadataMap insertion order and JSON object order on disk. + let mut value = serde_json::to_value(value).map_err(|error| error.to_string())?; + value.sort_all_objects(); + let bytes = serde_json::to_vec(&value).map_err(|error| error.to_string())?; + Ok(blake3::hash(&bytes).to_hex().to_string()) +} + +fn prefix_digest(prefix: &[Item]) -> Result { + let mut prefix = prefix.to_vec(); + if let Some(bootstrap) = prefix.first_mut() { + bootstrap.metadata.remove(METADATA_KEY); + } + digest(&prefix) +} + +pub(crate) fn load_history(root: &Path, session_id: &str) -> Result>, String> { + load_history_in(root, &default_directory()?, session_id) +} + +pub(crate) fn load_history_in( + root: &Path, + directory: &Path, + session_id: &str, +) -> Result>, String> { + validate_id(session_id)?; + let authority = select_authority(directory, &canonical_workspace(root), session_id)? + .ok_or_else(|| format!("session {session_id:?} does not exist"))?; + for state in &authority.historical_items { + BranchMetadata::read(state)?; + } + // Do not repair or stamp items: boundary provenance refers to exact stored states. + Ok(authority.historical_items) +} + +/// Commit only a newly opened, still-guarded branch. Appending a bootstrap (or +/// even all initial items) is not completion. Keep the creation guard armed until +/// the matching replacement has crossed the branch-scoped disk barrier. +pub(crate) fn commit(observer: &SessionObserver, transcript: &[Item]) -> Result<(), String> { + commit_with_barrier(observer, transcript, |filesystem, path| { + filesystem.require_disk(path) + }) +} + +fn commit_with_barrier( + observer: &SessionObserver, + transcript: &[Item], + barrier: impl FnOnce(&Fs, &Path) -> io::Result<()>, +) -> Result<(), String> { + let metadata = BranchMetadata::read(transcript)? + .ok_or_else(|| "missing prompt checkout metadata".to_string())?; + if !metadata.matches_snapshot(transcript)? { + return Err("incomplete or mismatched prompt checkout transcript".into()); + } + let mut writer = observer + .0 + .lock() + .map_err(|_| "session transcript writer poisoned".to_string())?; + if writer.created.is_none() { + return Err("prompt checkout commit requires an uncommitted new session".into()); + } + if metadata.parent_session_id == writer.session_id { + return Err("prompt checkout cannot parent itself".into()); + } + writer.ensure_lock()?; + let StoredTranscript::History(history) = read_records_direct(&writer.path, &writer.session_id)? + else { + return Err("prompt checkout transcript is a redirect".into()); + }; + if history.items != transcript { + return Err("prompt checkout initial transcript differs from commit snapshot".into()); + } + writer.replace(transcript)?; + barrier(&writer.lock.filesystem()?, &writer.path) + .map_err(|error| format!("prompt checkout is not durable: {error}"))?; + writer.commit_creation() +} + +#[derive(Clone, Debug)] +pub(crate) struct CommittedBranch { + pub session_id: String, + pub metadata: BranchMetadata, + /// The committed prefix + submitted prompt, not later assistant output. + pub transcript: Vec, +} + +/// Restart-safe idempotency lookup. None means absent, legacy/root, a different +/// request, or incomplete initialization. Corrupt/unknown metadata is an error. +/// Later appends and compaction do not erase an earlier commit replacement. +pub(crate) fn lookup_committed( + root: &Path, + session_id: &str, + checkout_id: &str, + request_id: &str, +) -> Result, String> { + lookup_committed_in( + root, + &default_directory()?, + session_id, + checkout_id, + request_id, + ) +} + +pub(crate) fn lookup_committed_in( + root: &Path, + directory: &Path, + session_id: &str, + checkout_id: &str, + request_id: &str, +) -> Result, String> { + lookup_committed_with_parent_in(root, directory, session_id, checkout_id, request_id, None) +} + +fn lookup_committed_with_parent_in( + root: &Path, + directory: &Path, + session_id: &str, + checkout_id: &str, + request_id: &str, + parent_session_id: Option<&str>, +) -> Result, String> { + validate_id(session_id)?; + let Some(authority) = select_authority_with( + directory, + &canonical_workspace(root), + session_id, + true, + true, + )? + else { + return Ok(None); + }; + // Discovery validates unrelated metadata but must not demand durability + // from unrelated sessions. Only candidates for this token cross a barrier. + let mut candidate = false; + for state in &authority.historical_items { + if let Some(metadata) = BranchMetadata::read(state)? + && metadata.checkout_id == checkout_id + { + candidate = true; + } + } + if !candidate { + return Ok(None); + } + // An overlay-only record cannot establish completion, even in this process. + fs::require_disk(&authority.path) + .map_err(|error| format!("prompt checkout is not durable: {error}"))?; + // A later interrupted append does not undo a complete replacement. Read a + // bounded snapshot after the disk barrier; only the locked opener may + // truncate its incomplete final record. Partial replacements cannot commit. + let bytes = transcript_snapshot(&authority.path, true)?; + let StoredTranscript::History(history) = + read_records_bytes(&authority.path, session_id, &bytes)? + else { + return Err("prompt checkout authority is a redirect".into()); + }; + for state in &history.states { + if let Some(metadata) = BranchMetadata::read(state)? + && metadata.checkout_id == checkout_id + && let Some(parent) = parent_session_id + && (metadata.parent_session_id != parent || metadata.request.id != request_id) + { + // An initialized but not yet committed child also binds this token. + // Check every state: compaction must not hide identity conflicts. + return Err( + "this checkout token is already bound to a different submitted request or source" + .into(), + ); + } + } + let mut found: Option = None; + for (state_index, length) in history.replacement_boundaries { + let snapshot = &history.states[state_index][..length]; + let Some(metadata) = BranchMetadata::read(snapshot)? else { + continue; + }; + if metadata.parent_session_id == session_id { + return Err("prompt checkout cannot parent itself".into()); + } + if metadata.checkout_id == checkout_id + && metadata.request.id == request_id + && metadata.matches_snapshot(snapshot)? + { + if found + .as_ref() + .is_some_and(|previous| previous.transcript != snapshot) + { + return Err("conflicting prompt checkout completion records".into()); + } + found = Some(CommittedBranch { + session_id: session_id.into(), + metadata, + transcript: snapshot.to_vec(), + }); + } + } + Ok(found) +} + +/// Discover a committed child when a process restart lost the destination ID. +/// New checkout branches are workspace-scoped; no sidecar/index is required. +/// Multiple durable children for one request are an explicit conflict. Reusing +/// a checkout token with another request or parent is an error, including when +/// its original child is only partially initialized. +pub(crate) fn find_committed( + root: &Path, + parent_session_id: &str, + checkout_id: &str, + request_id: &str, +) -> Result, String> { + find_committed_in( + root, + &default_directory()?, + parent_session_id, + checkout_id, + request_id, + ) +} + +fn find_committed_in( + root: &Path, + directory: &Path, + parent_session_id: &str, + checkout_id: &str, + request_id: &str, +) -> Result, String> { + validate_id(parent_session_id)?; + let scoped = workspace_storage_directory(directory, &canonical_workspace(root)); + let mut found = None; + for session_id in list_ids_in(&scoped)? { + if let Some(committed) = lookup_committed_with_parent_in( + root, + directory, + &session_id, + checkout_id, + request_id, + Some(parent_session_id), + )? { + if found.is_some() { + return Err("multiple committed prompt checkout children for one request".into()); + } + found = Some(committed); + } + } + Ok(found) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn prepared() -> Vec { + let prefix = vec![Item::text(ItemKind::System, "bootstrap")]; + prepare( + prefix.clone(), + "parent".into(), + Boundary::new(0, &prefix).unwrap(), + "checkout-1".into(), + SubmittedRequest { + id: "request-1".into(), + selection: CapturedSelection::new( + &ModelSelection::new(crate::ProviderKind::OpenRouter, "provider/model"), + Some(ReasoningEffort::High), + ), + }, + Item::text(ItemKind::User, "submitted prompt"), + ) + .unwrap() + } + + fn open_branch(root: &Path, initial: Vec) -> OpenSession { + open_with_initial_timestamps_in( + root, + &root.join("sessions"), + "branch", + false, + false, + initial, + InitialTranscriptOptions { + stamp_items: false, + commit_creation: false, + }, + ) + .unwrap() + } + + fn path(root: &Path) -> PathBuf { + transcript_path( + &workspace_storage_directory(&root.join("sessions"), &canonical_workspace(root)), + "branch", + ) + } + + fn lookup(root: &Path) -> Result, String> { + lookup_committed_in( + root, + &root.join("sessions"), + "branch", + "checkout-1", + "request-1", + ) + } + + #[test] + fn current_metadata_round_trips_and_missing_legacy_is_valid() { + let transcript = prepared(); + let metadata = BranchMetadata::read(&transcript).unwrap().unwrap(); + assert_eq!(metadata.version, 1); + assert_eq!(metadata.parent_session_id, "parent"); + assert!(metadata.matches_snapshot(&transcript).unwrap()); + assert_eq!( + metadata.request.selection.resolve().unwrap().1, + Some(ReasoningEffort::High) + ); + assert_eq!(transcript[0].created_at, None); + let mut legacy = transcript; + legacy[0].metadata.remove(METADATA_KEY); + assert!(BranchMetadata::read(&legacy).unwrap().is_none()); + assert!(BranchMetadata::read(&[]).unwrap().is_none()); + } + + #[test] + fn malformed_and_unknown_nested_payloads_are_explicit_errors() { + let transcript = prepared(); + let good = transcript[0].metadata[METADATA_KEY].clone(); + let mut misplaced = transcript.clone(); + misplaced[0].metadata.remove(METADATA_KEY); + misplaced[1] + .metadata + .insert(METADATA_KEY.into(), good.clone()); + assert!(BranchMetadata::read(&misplaced).is_err()); + let mut wrong_bootstrap = transcript.clone(); + wrong_bootstrap[0].kind = ItemKind::User; + assert!(BranchMetadata::read(&wrong_bootstrap).is_err()); + for value in [ + serde_json::Value::Null, + json!({}), + json!({"version": 2}), + json!({"version": "1"}), + ] { + let mut invalid = transcript.clone(); + invalid[0].metadata.insert(METADATA_KEY.into(), value); + assert!(BranchMetadata::read(&invalid).is_err()); + } + for pointer in [ + "", + "/boundary", + "/request", + "/request/selection", + "/completion", + ] { + let mut value = good.clone(); + value + .pointer_mut(pointer) + .unwrap() + .as_object_mut() + .unwrap() + .insert("unknown".into(), json!(true)); + let mut invalid = transcript.clone(); + invalid[0].metadata.insert(METADATA_KEY.into(), value); + assert!( + BranchMetadata::read(&invalid) + .unwrap_err() + .contains("unknown field") + ); + } + for pointer in [ + "/request/selection/reasoning", + "/request/selection/provider", + "/completion/prefix_hash", + ] { + let mut value = good.clone(); + *value.pointer_mut(pointer).unwrap() = json!("bogus"); + let mut invalid = transcript.clone(); + invalid[0].metadata.insert(METADATA_KEY.into(), value); + assert!(BranchMetadata::read(&invalid).is_err()); + } + let mut value = good; + value["request"]["selection"] + .as_object_mut() + .unwrap() + .remove("reasoning"); + let mut invalid = transcript; + invalid[0].metadata.insert(METADATA_KEY.into(), value); + assert!(BranchMetadata::read(&invalid).is_err()); + } + + #[test] + fn descendants_replace_inherited_payload_and_preserve_other_metadata() { + let mut prefix = prepared(); + prefix[0].metadata.insert("unrelated".into(), json!(true)); + let boundary = Boundary::new(3, &prefix).unwrap(); + let request = SubmittedRequest { + id: "descendant-request".into(), + selection: CapturedSelection::new( + &ModelSelection::new(crate::ProviderKind::Speakeasy, "model"), + None, + ), + }; + let descendant = prepare( + prefix, + "branch".into(), + boundary.clone(), + "descendant-checkout".into(), + request, + Item::text(ItemKind::User, "next"), + ) + .unwrap(); + let metadata = BranchMetadata::read(&descendant).unwrap().unwrap(); + assert_eq!(metadata.parent_session_id, "branch"); + assert_eq!(metadata.boundary, boundary); + assert_eq!(metadata.checkout_id, "descendant-checkout"); + assert_eq!(metadata.request.id, "descendant-request"); + assert_eq!(metadata.request.selection.reasoning, "default"); + assert_eq!(descendant[0].metadata["unrelated"], json!(true)); + assert!(metadata.matches_snapshot(&descendant).unwrap()); + } + + #[test] + fn writer_uses_existing_replacement_shape_and_restart_lookup_survives_compaction() { + let root = tempfile::tempdir().unwrap(); + let transcript = prepared(); + let opened = open_branch(root.path(), transcript.clone()); + assert!( + lookup(root.path()).unwrap().is_none(), + "initial appends are not a commit" + ); + commit(&opened.observer, &opened.transcript).unwrap(); + let bytes = std::fs::read(path(root.path())).unwrap(); + let records: Vec = std::str::from_utf8(&bytes) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + assert!( + records + .iter() + .all(|record| record["schema_version"] == SCHEMA_VERSION) + ); + let last = records.last().unwrap(); + assert!(last.get("replacement").is_some()); + assert!(last.get("item").is_none()); + assert_eq!(last.as_object().unwrap().len(), 5); + assert_eq!( + last["replacement"][0]["metadata"][METADATA_KEY]["version"], + 1 + ); + opened + .observer + .replace(&[Item::text(ItemKind::System, "compacted")]) + .unwrap(); + drop(opened); + let before = std::fs::read(path(root.path())).unwrap(); + let committed = lookup(root.path()).unwrap().unwrap(); + assert_eq!(committed.transcript, transcript); + assert_eq!(committed.session_id, "branch"); + assert_eq!(committed.metadata.request.id, "request-1"); + assert_eq!( + std::fs::read(path(root.path())).unwrap(), + before, + "lookup must not rewrite" + ); + assert!( + lookup_committed_in( + root.path(), + &root.path().join("sessions"), + "branch", + "checkout-1", + "other" + ) + .unwrap() + .is_none() + ); + } + + #[test] + fn committed_lookup_ignores_only_torn_final_appends_and_opener_recovers() { + for kind in [ItemKind::Assistant, ItemKind::Tool] { + let root = tempfile::tempdir().unwrap(); + let transcript = prepared(); + let opened = open_branch(root.path(), transcript.clone()); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let path = path(root.path()); + let complete = std::fs::read(&path).unwrap(); + let append = Record { + schema_version: SCHEMA_VERSION, + session_id: "branch".into(), + generation: 4, + workspace_root: Some(canonical_workspace(root.path())), + item: Some(Item::text(kind, "later output")), + replacement: None, + redirect: None, + }; + let encoded = serde_json::to_vec(&append).unwrap(); + let mut torn = complete.clone(); + torn.extend_from_slice(&encoded[..encoded.len() - 2]); + std::fs::write(&path, &torn).unwrap(); + for _ in 0..2 { + let recovered = lookup(root.path()).unwrap().unwrap(); + assert_eq!(recovered.session_id, "branch"); + assert_eq!(recovered.transcript, transcript); + let discovered = find_committed_in( + root.path(), + &root.path().join("sessions"), + "parent", + "checkout-1", + "request-1", + ) + .unwrap() + .unwrap(); + assert_eq!(discovered.session_id, recovered.session_id); + assert_eq!(discovered.metadata, recovered.metadata); + assert_eq!(discovered.transcript, recovered.transcript); + assert_eq!( + std::fs::read(&path).unwrap(), + torn, + "lookup must not repair" + ); + } + let opened = open_in( + root.path(), + &root.path().join("sessions"), + "branch", + true, + false, + Vec::new(), + ) + .unwrap(); + assert_eq!(opened.transcript, transcript); + drop(opened); + assert_eq!( + std::fs::read(&path).unwrap(), + complete, + "only the locked opener removes the torn append, without a generation rerun" + ); + } + } + + #[test] + fn committed_lookup_and_opener_reject_malformed_complete_records() { + for tail in [ + b"{}".as_slice(), + b"{}\n", + b"{not json}", + b"{not json}\n", + b"{}\n{\"schema_version\":", + b"{not json\n{\"schema_version\":", + b"{not json\xe2", + b"{\"schema_version\":\n", + ] { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let path = path(root.path()); + let mut bytes = std::fs::read(&path).unwrap(); + bytes.extend_from_slice(tail); + std::fs::write(&path, &bytes).unwrap(); + assert!(lookup(root.path()).is_err(), "tail: {tail:?}"); + assert_eq!(std::fs::read(&path).unwrap(), bytes); + assert!( + open_in( + root.path(), + &root.path().join("sessions"), + "branch", + true, + false, + Vec::new() + ) + .is_err(), + "tail: {tail:?}" + ); + // Recovery must not erase a malformed complete final record. + if !tail.windows(2).any(|pair| pair == b"\n{") { + assert_eq!(std::fs::read(&path).unwrap(), bytes); + } + } + } + + #[test] + fn torn_initial_or_completion_record_does_not_establish_completion() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + let initial = std::fs::read(path(root.path())).unwrap(); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let complete = std::fs::read(path(root.path())).unwrap(); + // Exercise every byte boundary, including an empty initial record. + // No partial replacement can prove the commit. + for length in 0..complete.len() - 1 { + std::fs::write(path(root.path()), &complete[..length]).unwrap(); + assert!( + !matches!(lookup(root.path()), Ok(Some(_))), + "length {length}" + ); + assert_eq!( + std::fs::read(path(root.path())).unwrap(), + complete[..length] + ); + } + assert!(initial.len() < complete.len() - 2); + // A fully present replacement is valid even if only its newline tore. + std::fs::write(path(root.path()), &complete[..complete.len() - 1]).unwrap(); + assert!(lookup(root.path()).unwrap().is_some()); + } + + #[test] + fn newline_only_torn_completion_survives_resume_append_and_reload() { + let root = tempfile::tempdir().unwrap(); + let transcript = prepared(); + let opened = open_branch(root.path(), transcript.clone()); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let path = path(root.path()); + let complete = std::fs::read(&path).unwrap(); + assert_eq!(complete.last(), Some(&b'\n')); + let unterminated = &complete[..complete.len() - 1]; + std::fs::write(&path, unterminated).unwrap(); + assert_eq!(lookup(root.path()).unwrap().unwrap().transcript, transcript); + assert_eq!(std::fs::read(&path).unwrap(), unterminated); + + let resumed = open_in( + root.path(), + &root.path().join("sessions"), + "branch", + true, + false, + Vec::new(), + ) + .unwrap(); + assert_eq!(resumed.transcript, transcript); + assert_eq!(std::fs::read(&path).unwrap(), complete); + let appended = Item::text(ItemKind::Assistant, "after resume") + .with_created_at(agentkit_core::Timestamp(123)); + resumed.observer.on_transcript_event(TranscriptEvent { + session_id: &agentkit_core::SessionId::new("branch"), + item: &appended, + }); + drop(resumed); + + let mut expected = transcript.clone(); + expected.push(appended); + assert_eq!( + load_in(root.path(), &root.path().join("sessions"), "branch").unwrap(), + expected + ); + let after = std::fs::read(&path).unwrap(); + assert!(after.starts_with(&complete)); + assert_eq!( + after.iter().filter(|byte| **byte == b'\n').count(), + complete.iter().filter(|byte| **byte == b'\n').count() + 1 + ); + assert_eq!(lookup(root.path()).unwrap().unwrap().transcript, transcript); + assert_eq!(std::fs::read(&path).unwrap(), after); + } + + #[test] + fn history_includes_precompaction_and_reads_do_not_repair_or_rewrite() { + let root = tempfile::tempdir().unwrap(); + let initial = vec![ + Item::text(ItemKind::System, "bootstrap"), + Item::text(ItemKind::User, "old"), + ]; + let opened = open_branch(root.path(), initial.clone()); + opened.observer.commit_creation().unwrap(); + let compacted = vec![Item::text(ItemKind::System, "compacted")]; + opened.observer.replace(&compacted).unwrap(); + drop(opened); + let before = std::fs::read(path(root.path())).unwrap(); + assert_eq!( + load_history_in(root.path(), &root.path().join("sessions"), "branch").unwrap(), + vec![initial, compacted] + ); + assert_eq!(std::fs::read(path(root.path())).unwrap(), before); + } + + #[test] + fn partial_initialization_is_not_completion_and_guard_cleans_up() { + for count in 1..=2 { + let root = tempfile::tempdir().unwrap(); + let full = prepared(); + let opened = open_branch(root.path(), full[..count].to_vec()); + assert!(lookup(root.path()).unwrap().is_none()); + if count == 1 { + assert!(commit(&opened.observer, &opened.transcript).is_err()); + assert!(commit(&opened.observer, &full).is_err()); + } + drop(opened); + assert!(!path(root.path()).exists()); + } + } + + #[test] + fn wrong_prefix_prompt_or_length_cannot_commit_or_recover() { + for mutation in 0..3 { + let root = tempfile::tempdir().unwrap(); + let mut transcript = prepared(); + match mutation { + 0 => transcript[0].parts = Item::text(ItemKind::System, "wrong bootstrap").parts, + 1 => transcript[1].parts = Item::text(ItemKind::User, "wrong prompt").parts, + _ => transcript.push(Item::text(ItemKind::User, "extra")), + } + let opened = open_branch(root.path(), transcript.clone()); + assert!(commit(&opened.observer, &transcript).is_err()); + // Even a well-formed replacement is insufficient without full content validation. + opened.observer.replace(&transcript).unwrap(); + assert!(lookup(root.path()).unwrap().is_none()); + } + } + + #[test] + fn failed_branch_scoped_barrier_keeps_creation_guard_armed() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + let error = commit_with_barrier(&opened.observer, &opened.transcript, |_, actual_path| { + assert_eq!(actual_path, path(root.path())); + Err(io::Error::from_raw_os_error(libc::ENOSPC)) + }) + .unwrap_err(); + assert!(error.contains("not durable")); + assert!(opened.observer.0.lock().unwrap().created.is_some()); + drop(opened); + assert!(!path(root.path()).exists()); + assert!(lookup(root.path()).unwrap().is_none()); + } + + #[test] + fn failed_barrier_can_retry_without_duplicate_recovery_results() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + assert!( + commit_with_barrier(&opened.observer, &opened.transcript, |_, _| { + Err(io::Error::from_raw_os_error(libc::ENOSPC)) + }) + .is_err() + ); + commit(&opened.observer, &opened.transcript).unwrap(); + assert!(opened.observer.0.lock().unwrap().created.is_none()); + let expected = opened.transcript.clone(); + drop(opened); + assert_eq!(lookup(root.path()).unwrap().unwrap().transcript, expected); + } + + #[test] + fn restart_discovery_finds_request_without_destination_id() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + assert!( + find_committed_in( + root.path(), + &root.path().join("sessions"), + "parent", + "checkout-1", + "request-1" + ) + .unwrap() + .is_none() + ); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let found = find_committed_in( + root.path(), + &root.path().join("sessions"), + "parent", + "checkout-1", + "request-1", + ) + .unwrap() + .unwrap(); + assert_eq!(found.session_id, "branch"); + assert!( + find_committed_in( + root.path(), + &root.path().join("sessions"), + "other-parent", + "checkout-1", + "request-1" + ) + .is_err() + ); + } + + #[test] + fn discovery_rejects_request_and_source_mismatches_even_before_commit() { + for stage in 0..3 { + let root = tempfile::tempdir().unwrap(); + let full = prepared(); + let initial = if stage == 0 { full[..1].to_vec() } else { full }; + let opened = open_branch(root.path(), initial); + if stage == 2 { + commit(&opened.observer, &opened.transcript).unwrap(); + // The conflicting identity must remain visible in history even + // if a later replacement has removed the current payload. + opened + .observer + .replace(&[Item::text(ItemKind::System, "compacted")]) + .unwrap(); + } + for (parent, request) in [ + ("parent", "different-request"), + ("different-parent", "request-1"), + ] { + let error = find_committed_in( + root.path(), + &root.path().join("sessions"), + parent, + "checkout-1", + request, + ) + .unwrap_err(); + assert!(error.contains("different submitted request or source")); + } + assert!( + find_committed_in( + root.path(), + &root.path().join("sessions"), + "parent", + "unrelated-checkout", + "request-1" + ) + .unwrap() + .is_none() + ); + } + } + + #[test] + fn misplaced_duplicate_malformed_and_unknown_payloads_never_look_legacy() { + let original = prepared(); + for value in [ + original[0].metadata[METADATA_KEY].clone(), + serde_json::Value::Null, + json!({"version": 999}), + ] { + for keep_bootstrap in [false, true] { + let mut transcript = original.clone(); + if !keep_bootstrap { + transcript[0].metadata.remove(METADATA_KEY); + } + transcript[1] + .metadata + .insert(METADATA_KEY.into(), value.clone()); + assert!( + BranchMetadata::read(&transcript) + .unwrap_err() + .contains("bootstrap") + ); + } + let mut transcript = original.clone(); + transcript[0].kind = ItemKind::User; + transcript[0].metadata.insert(METADATA_KEY.into(), value); + assert!( + BranchMetadata::read(&transcript) + .unwrap_err() + .contains("system bootstrap") + ); + } + assert!( + BranchMetadata::read(&[Item::text(ItemKind::User, "legacy without bootstrap")]) + .unwrap() + .is_none() + ); + } + + #[test] + fn later_appends_cannot_complete_a_short_replacement_record() { + let root = tempfile::tempdir().unwrap(); + let full = prepared(); + let opened = open_branch(root.path(), full[..1].to_vec()); + opened.observer.replace(&full[..1]).unwrap(); + opened.observer.0.lock().unwrap().append(&full[1]).unwrap(); + // Canonical history now equals the intended branch. But no single + // replacement record contains its entire prefix and submitted prompt. + assert_eq!(read_records(&path(root.path()), "branch").unwrap().0, full); + assert!(lookup(root.path()).unwrap().is_none()); + opened.observer.commit_creation().unwrap(); // Model a surviving partial initialization. + drop(opened); + assert!(lookup(root.path()).unwrap().is_none()); + } + + #[test] + fn legacy_history_is_read_without_stamp_or_schema_rewrite() { + let root = tempfile::tempdir().unwrap(); + let initial = vec![ + Item::text(ItemKind::System, "old bootstrap"), + Item::text(ItemKind::User, "old prompt"), + ]; + let opened = open_branch(root.path(), initial.clone()); + opened.observer.commit_creation().unwrap(); + let compacted = vec![Item::text(ItemKind::System, "old compaction")]; + opened.observer.replace(&compacted).unwrap(); + drop(opened); + let file = path(root.path()); + let mut records: Vec = std::fs::read_to_string(&file) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + for record in &mut records { + record["schema_version"] = json!(if record.get("replacement").is_some() { + PREVIOUS_SCHEMA_VERSION + } else { + LEGACY_SCHEMA_VERSION + }); + } + let bytes = records + .iter() + .map(|record| format!("{record}\n")) + .collect::(); + std::fs::write(&file, &bytes).unwrap(); + assert_eq!( + load_history_in(root.path(), &root.path().join("sessions"), "branch").unwrap(), + vec![initial, compacted] + ); + assert_eq!(std::fs::read_to_string(&file).unwrap(), bytes); + assert!(lookup(root.path()).unwrap().is_none()); + } + + #[test] + fn lookup_rejects_corrupt_generation_and_does_not_accept_torn_replacement() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), prepared()); + commit(&opened.observer, &opened.transcript).unwrap(); + drop(opened); + let file = path(root.path()); + let original = std::fs::read(&file).unwrap(); + std::fs::write(&file, &original[..original.len() - 12]).unwrap(); + assert!(lookup(root.path()).unwrap().is_none()); + assert_eq!( + std::fs::read(&file).unwrap(), + original[..original.len() - 12] + ); + let mut lines: Vec = std::str::from_utf8(&original) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + lines[1]["generation"] = json!(999); + std::fs::write( + &file, + lines + .iter() + .map(|line| format!("{line}\n")) + .collect::(), + ) + .unwrap(); + assert!(lookup(root.path()).is_err()); + } +} diff --git a/src/tools/mcp.rs b/src/tools/mcp.rs index 9338a58..d9186db 100644 --- a/src/tools/mcp.rs +++ b/src/tools/mcp.rs @@ -405,6 +405,10 @@ pub(crate) struct McpSubscription { } impl McpSubscription { + pub(crate) fn has_pending(&self) -> bool { + !self.receiver.is_empty() + } + pub(crate) async fn recv(&mut self) -> Option { self.receiver.recv().await } diff --git a/src/tui/app.rs b/src/tui/app.rs index a5f98c2..e5b7186 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -118,6 +118,11 @@ pub enum Update { State(StateUpdate), /// A nested tool call started or finished inside a compose run. Runtime(RuntimeEvent), + /// Session identity captured from the ordered stderr stream, before queuing. + RoutedRuntime { + session_id: String, + event: RuntimeEvent, + }, /// A diagnostic line from the agent process. Log(String), /// The ACP process exited while work could still be active. @@ -276,7 +281,36 @@ struct SteerEdit { next_attachment: usize, } +pub(super) const BRANCH_WARNING: &str = "Only conversation context changes. Filesystem changes, running processes, and external effects are not rolled back."; + +pub(super) struct BranchChooser { + pub boundaries: Vec, + pub selected: usize, + pub pending: bool, +} + +struct BranchDraft { + original: Box, + checkout_token: String, + submitting: bool, + // stderr can reach the UI before the submit response installs the child. + // Keep its session attribution until canonical ACP replay is applied. + early_runtime: Vec<(String, RuntimeEvent)>, +} + pub enum Action { + ListPromptBranches { + epoch: u64, + }, + PreparePromptBranch { + epoch: u64, + address: String, + }, + SubmitPromptBranch { + epoch: u64, + checkout_token: String, + text: String, + }, None, Redraw, Submit { @@ -612,6 +646,12 @@ pub struct AgentCounts { } pub struct App { + pub(super) branch_epoch: u64, + pub(super) branch_chooser: Option, + branch_draft: Option, + // AvailableCommands is emitted after canonical branch replay. stderr may + // outrun that ordered ACP stream, so retain diagnostics until it arrives. + branch_runtime_replay: Option>, pub root: PathBuf, pub provider: String, pub model: String, @@ -879,6 +919,10 @@ impl App { session_dialog: None, file_picker: None, session_catalog_pending: false, + branch_epoch: 0, + branch_chooser: None, + branch_draft: None, + branch_runtime_replay: None, auth_methods: Vec::new(), available_commands: Vec::new(), command_completion_selected: 0, @@ -960,7 +1004,11 @@ impl App { } fn command_completion_prefix(&self) -> Option<&str> { - if self.editing_steer() || self.queue_focused { + if self.editing_steer() + || self.queue_focused + || self.editing_branch() + || self.branch_chooser.is_some() + { return None; } command::completion_prefix(self.editor.text(), self.editor.cursor()) @@ -1607,6 +1655,65 @@ impl App { } pub fn apply(&mut self, update: Update) { + // The source stays loaded while its provisional replacement is visible. + if let Some(draft) = &mut self.branch_draft { + match update { + Update::RoutedRuntime { session_id, event } => { + if draft.original.session_id.as_ref() == Some(&session_id) { + draft + .original + .apply(Update::RoutedRuntime { session_id, event }); + } else { + draft.early_runtime.push((session_id, event)); + } + } + Update::Runtime(RuntimeEvent::StorageStatus { pending, exhausted }) => { + // Durability is process-global, not part of the parked view. + self.storage_pending = pending; + self.storage_exhausted = exhausted; + draft + .original + .apply(Update::Runtime(RuntimeEvent::StorageStatus { + pending, + exhausted, + })); + } + Update::Runtime(RuntimeEvent::SessionStarted { session_id }) => { + self.runtime_session_id = Some(session_id.clone()); + if draft.original.session_id.as_ref() == Some(&session_id) { + draft.original.activate_runtime_session(); + } + } + Update::Runtime(event) => { + if self.runtime_session_id == draft.original.session_id { + draft.original.apply(Update::Runtime(event)); + } else if let Some(session_id) = &self.runtime_session_id { + draft.early_runtime.push((session_id.clone(), event)); + } + } + update => { + if let Update::ConfigOptions(options) = &update { + super::refresh_config_state(&mut draft.original, Some(options)); + } + draft.original.apply(update); + } + } + return; + } + if let Some(pending) = &mut self.branch_runtime_replay + && let Update::Runtime(event) = &update + && !matches!( + event, + RuntimeEvent::StorageStatus { .. } | RuntimeEvent::SessionStarted { .. } + ) + { + if self.runtime_session_id == self.session_id { + pending.push(event.clone()); + } + return; + } + let replay_complete = matches!(&update, Update::AvailableCommands { session_id, .. } + if self.session_id.as_ref() == Some(session_id)); match update { Update::A2aAddress(address) => self.a2a = address, Update::SessionCatalog(result) => { @@ -1821,6 +1928,14 @@ impl App { Update::Usage { used, size } => { self.usage = Some(ContextUsage { used, size }); } + Update::RoutedRuntime { session_id, event } => { + if self.session_id.as_ref() == Some(&session_id) { + // A queued marker from an older activation cannot change this + // event's identity or the route established by ACP activation. + self.activate_runtime_session(); + self.apply(Update::Runtime(event)); + } + } Update::Runtime(event) => self.apply_runtime(event), Update::Log(line) => { self.logs.push(line); @@ -1854,6 +1969,16 @@ impl App { self.push_block(Block::Error(error)); } } + if replay_complete && let Some(events) = self.branch_runtime_replay.take() { + // These events were attributed to the active child when queued. + let route = self + .runtime_session_id + .replace(self.session_id.clone().unwrap()); + for event in events { + self.apply_runtime(event); + } + self.runtime_session_id = route; + } if self.follow { self.scroll = usize::MAX; } @@ -1971,10 +2096,19 @@ impl App { } } + /// A successful ACP activation establishes the diagnostic route explicitly. + /// Stderr markers only label ingress events; ACP owns the visible route. + pub(super) fn activate_runtime_session(&mut self) { + self.runtime_session_id = self.session_id.clone(); + } + /// Switches the visible client state to a fresh persisted session. Editor /// history and diagnostics remain useful, while transcript-derived state /// starts empty. pub fn start_session(&mut self, session_id: String) { + self.abandon_branch(); + self.branch_runtime_replay = None; + self.branch_epoch = self.branch_epoch.wrapping_add(1); self.model_switch = None; self.cancel_steer_edit(); self.selected_steer = None; @@ -2216,6 +2350,10 @@ impl App { kind: AttachmentKind, size: u64, ) { + if self.branch_draft.is_some() || self.branch_chooser.is_some() { + self.toast("prompt branch edits are text-only"); + return; + } if self.editing_steer() { self.toast("pending-message edits are text-only"); return; @@ -2288,6 +2426,14 @@ impl App { if self.model_switch.is_some() { return; } + if self.branch_chooser.is_some() || self.branch_submitting() { + return; + } + if self.branch_draft.is_some() { + self.last_key = None; + self.editor.insert_str(text); + return; + } if let Some(dialog) = &mut self.navigation.dialog { dialog.insert(text); self.sync_navigation(); @@ -2940,6 +3086,16 @@ impl App { .and_then(|dialog| dialog.selected); let current = selected.and_then(|id| self.navigation.index(id)); match key.code { + KeyCode::Enter + if key.modifiers.is_empty() + && self + .navigation + .dialog + .as_ref() + .is_some_and(|dialog| dialog.query == "/branch") => + { + return self.open_branch_chooser(); + } KeyCode::Esc | KeyCode::F(3) => self.navigation.dialog = None, KeyCode::Enter if key.modifiers.is_empty() => { if let Some(index) = current { @@ -3050,6 +3206,254 @@ impl App { self.navigation.reveal_pending = false; } + /// Applies a key press, returning work for the event loop. + pub(super) fn editing_branch(&self) -> bool { + self.branch_draft.is_some() + } + + pub(super) fn branch_submitting(&self) -> bool { + self.branch_draft + .as_ref() + .is_some_and(|draft| draft.submitting) + } + + fn open_branch_chooser(&mut self) -> Action { + if self.working() || self.editing_steer() || !self.pending_steers.is_empty() { + self.toast( + "prompt checkout is available only while idle and outside pending-message edits", + ); + return Action::None; + } + self.branch_epoch = self.branch_epoch.wrapping_add(1); + self.branch_chooser = Some(BranchChooser { + boundaries: Vec::new(), + selected: 0, + pending: true, + }); + Action::ListPromptBranches { + epoch: self.branch_epoch, + } + } + + pub(super) fn branch_listed( + &mut self, + epoch: u64, + result: Result, String>, + ) { + if epoch != self.branch_epoch { + return; + } + let Some(chooser) = &mut self.branch_chooser else { + return; + }; + match result { + Ok(boundaries) => { + chooser.boundaries = boundaries; + chooser.pending = false; + } + Err(error) => { + self.branch_chooser = None; + self.toast(format!("could not list prompt checkouts: {error}")); + } + } + } + + pub(super) fn branch_prepared( + &mut self, + epoch: u64, + result: Result, + ) { + if epoch != self.branch_epoch || self.branch_chooser.is_none() { + return; + } + let response = match result { + Ok(response) => response, + Err(error) => { + self.branch_chooser.as_mut().unwrap().pending = false; + self.toast(format!("could not prepare prompt checkout: {error}")); + return; + } + }; + let mut provisional = App::new( + self.root.clone(), + self.provider.clone(), + self.model.clone(), + self.a2a.clone(), + ); + provisional.session_id = self.session_id.clone(); + provisional.runtime_session_id = self.runtime_session_id.clone(); + provisional.storage_pending = self.storage_pending; + provisional.storage_exhausted = self.storage_exhausted; + provisional.last_key = self.last_key; + provisional.input_overflow = self.input_overflow; + provisional.input_recovery_ready = self.input_recovery_ready; + provisional.can_steer = self.can_steer; + provisional.can_replace_steer = self.can_replace_steer; + provisional.auth_methods = self.auth_methods.clone(); + provisional.show_thoughts = self.show_thoughts; + provisional.branch_epoch = epoch; + for update in response.prefix { + let (_, updates) = super::translate( + agent_client_protocol::schema::v2::UpdateSessionNotification::new( + self.session_id.clone().unwrap_or_default(), + update, + ), + ); + for update in updates { + provisional.apply(update); + } + } + super::refresh_config_state(&mut provisional, Some(&response.config_options)); + provisional.editor.insert_str(&response.original_text); + self.branch_chooser = None; + let original = Box::new(std::mem::replace(self, provisional)); + self.branch_draft = Some(BranchDraft { + original, + checkout_token: response.checkout_token, + submitting: false, + early_runtime: Vec::new(), + }); + } + + pub(super) fn abandon_branch(&mut self) { + let epoch = self.branch_epoch.wrapping_add(1); + if let Some(draft) = self.branch_draft.take() { + let input_state = ( + self.input_overflow, + self.input_recovery_ready, + self.last_key, + ); + *self = *draft.original; + ( + self.input_overflow, + self.input_recovery_ready, + self.last_key, + ) = input_state; + } + self.branch_chooser = None; + self.branch_epoch = epoch; + } + + pub(super) fn branch_submit_failed(&mut self, epoch: u64, error: String) { + if epoch != self.branch_epoch { + return; + } + if let Some(draft) = &mut self.branch_draft { + draft.submitting = false; + self.toast(format!("could not submit prompt checkout: {error}")); + } + } + + pub(super) fn branch_submitted(&mut self, epoch: u64, session_id: String) -> bool { + if epoch != self.branch_epoch || !self.branch_submitting() { + return false; + } + // Drop only the parked view, never close the source ACP session. stderr + // may precede both the response and canonical replay on the ACP stream. + let draft = self.branch_draft.take().expect("checked submitting draft"); + let early_runtime = draft + .early_runtime + .into_iter() + .filter_map(|(id, event)| (id == session_id).then_some(event)) + .collect(); + self.start_session(session_id); + self.activate_runtime_session(); + self.branch_runtime_replay = Some(early_runtime); + self.editor.clear(); + true + } + + fn handle_branch_key_at(&mut self, key: KeyEvent, arrival: Instant) -> Action { + let pasted = self + .last_key + .is_some_and(|last| arrival.saturating_duration_since(last) < PASTE_GAP); + self.handle_branch_key(key, pasted) + } + + fn handle_branch_key(&mut self, key: KeyEvent, pasted: bool) -> Action { + if self.branch_submitting() { + self.toast("creating branch — wait for the result"); + return Action::None; + } + if key.code == KeyCode::Esc { + self.abandon_branch(); + return Action::None; + } + if let Some(chooser) = &mut self.branch_chooser { + if chooser.pending { + return Action::None; + } + match key.code { + KeyCode::Up => chooser.selected = chooser.selected.saturating_sub(1), + KeyCode::Down => { + chooser.selected = + (chooser.selected + 1).min(chooser.boundaries.len().saturating_sub(1)) + } + KeyCode::Enter if key.modifiers.is_empty() && !pasted => { + if let Some(boundary) = chooser.boundaries.get(chooser.selected) { + let address = boundary.address.clone(); + chooser.pending = true; + self.branch_epoch = self.branch_epoch.wrapping_add(1); + return Action::PreparePromptBranch { + epoch: self.branch_epoch, + address, + }; + } + } + _ => {} + } + return Action::None; + } + let control = key.modifiers.contains(KeyModifiers::CONTROL); + match key.code { + KeyCode::Enter if key.modifiers.is_empty() && !pasted => { + if !self.attachments.is_empty() { + self.toast("prompt branch edits are text-only"); + return Action::None; + } + if self.editor.text().trim().is_empty() { + return Action::None; + } + let draft = self.branch_draft.as_mut().unwrap(); + draft.submitting = true; + self.branch_epoch = self.branch_epoch.wrapping_add(1); + return Action::SubmitPromptBranch { + epoch: self.branch_epoch, + checkout_token: draft.checkout_token.clone(), + text: self.editor.text().to_string(), + }; + } + KeyCode::Enter => self.editor.insert_char('\n'), + KeyCode::Char('j') if control => self.editor.insert_char('\n'), + KeyCode::Char('a') if control => self.editor.move_line_start(), + KeyCode::Char('e') if control => self.editor.move_line_end(), + KeyCode::Char('u') if control => self.editor.delete_to_line_start(), + KeyCode::Char('k') if control => self.editor.delete_to_line_end(), + KeyCode::Char('w') if control => self.editor.delete_word_back(), + KeyCode::Backspace if control => self.editor.delete_word_back(), + KeyCode::Backspace => self.editor.backspace(), + KeyCode::Delete => self.editor.delete_forward(), + KeyCode::Left => self.editor.move_left(), + KeyCode::Right => self.editor.move_right(), + KeyCode::Home => self.editor.move_line_start(), + KeyCode::End => self.editor.move_line_end(), + KeyCode::Up => { + self.editor.move_row_up(self.prompt_width); + } + KeyCode::Down => { + self.editor.move_row_down(self.prompt_width); + } + KeyCode::PageUp => self.scroll_by(-(self.viewport.max(2) as isize - 1)), + KeyCode::PageDown => self.scroll_by(self.viewport.max(2) as isize - 1), + KeyCode::Tab => self.editor.insert_str(" "), + KeyCode::Char(c) if !control && !key.modifiers.contains(KeyModifiers::SUPER) => { + self.editor.insert_char(c) + } + _ => {} + } + Action::None + } + /// Applies a key press using its terminal receipt time, not dispatch time. /// Synchronous rendering/search must not change inter-key paste gaps. pub fn handle_key_at(&mut self, key: KeyEvent, received_at: Instant) -> Action { @@ -3091,6 +3495,11 @@ impl App { } return Action::Redraw; } + if self.branch_chooser.is_some() || self.branch_draft.is_some() { + let action = self.handle_branch_key_at(key, received_at); + self.last_key = Some(received_at); + return action; + } if self.navigation.dialog.is_some() { let pasted = self .last_key @@ -3328,6 +3737,12 @@ impl App { self.open_navigation(); return Action::None; } + if matches!( + parse(self.editor.text(), !self.auth_methods.is_empty()), + Parsed::Branch + ) { + return self.open_branch_chooser(); + } if self.editor.is_empty() { return Action::None; } @@ -3399,6 +3814,7 @@ impl App { Action::ListSessions } } + Parsed::Branch => self.open_branch_chooser(), Parsed::Transcript => { self.open_navigation(); Action::None @@ -3579,7 +3995,11 @@ impl App { } pub fn handle_mouse(&mut self, mouse: MouseEvent) -> Action { - if self.model_switch.is_some() || self.navigation.dialog.is_some() { + if self.model_switch.is_some() + || self.navigation.dialog.is_some() + || self.branch_chooser.is_some() + || self.editing_branch() + { return Action::None; } match mouse.kind { @@ -4021,6 +4441,505 @@ mod tests { ) } + fn apply_stderr_runtime(app: &mut App, route: &mut Option, event: RuntimeEvent) { + let line = format!( + "{}{}", + crate::events::EVENT_MARKER, + serde_json::to_string(&event).unwrap() + ); + if let Some(update) = + super::super::runtime_diagnostic_update(route, crate::events::parse(&line).unwrap()) + { + app.apply(update); + } + } + + fn prepare_branch(app: &mut App) { + let Action::ListPromptBranches { epoch } = app.open_branch_chooser() else { + panic!("expected chooser") + }; + app.branch_prepared( + epoch, + Ok( + crate::protocols::acp::prompt_branches::PreparePromptBranchResponse { + checkout_token: "checkout-one".into(), + original_text: "original prompt".into(), + prefix: Vec::new(), + config_options: Vec::new(), + }, + ), + ); + } + + fn ready_branch_chooser(app: &mut App) -> u64 { + let Action::ListPromptBranches { epoch } = app.open_branch_chooser() else { + panic!("chooser") + }; + app.branch_listed( + epoch, + Ok(vec![ + crate::protocols::acp::prompt_branches::PromptBoundary { + address: "opaque".into(), + text: "original".into(), + historical: false, + }, + ]), + ); + epoch + } + + #[test] + fn branch_chooser_and_draft_exclude_slow_processing_but_preserve_real_input_gaps() { + // Dispatch time is irrelevant: use the terminal receipt time even when + // rendering/updating delayed dispatch of a buffered Enter. + for draft in [false, true] { + for deliberate in [false, true] { + let mut app = app(); + app.start_session("source".into()); + if draft { + prepare_branch(&mut app); + } else { + ready_branch_chooser(&mut app); + } + assert!(app.navigation.dialog.is_none()); // direct /branch path + let before = Instant::now(); + app.last_key = Some(if deliberate { + before - super::PASTE_GAP * 2 + } else { + before + }); + let action = app.handle_key_at(press(KeyCode::Enter), before); + if deliberate { + assert!(matches!(action, Action::SubmitPromptBranch { .. }) == draft); + assert!(matches!(action, Action::PreparePromptBranch { .. }) != draft); + } else { + assert!(matches!(action, Action::None)); + assert!(!app.branch_submitting()); + if draft { + assert_eq!(app.editor.text(), "original prompt\n"); + } else { + assert!(!app.branch_chooser.as_ref().unwrap().pending); + } + } + } + } + } + + #[test] + fn branch_prepare_response_preserves_paste_clock_in_the_fresh_provisional_view() { + let mut app = app(); + app.start_session("source".into()); + let epoch = ready_branch_chooser(&mut app); + let before = Instant::now(); + app.last_key = Some(before); + app.branch_prepared( + epoch, + Ok( + crate::protocols::acp::prompt_branches::PreparePromptBranchResponse { + checkout_token: "checkout".into(), + original_text: "original".into(), + prefix: Vec::new(), + config_options: Vec::new(), + }, + ), + ); + assert!(app.editing_branch()); + assert_eq!(app.last_key, Some(before)); + assert!(matches!( + app.handle_key_at(press(KeyCode::Enter), before), + Action::None + )); + assert!(!app.branch_submitting()); + assert_eq!(app.editor.text(), "original\n"); + } + + #[test] + fn branch_checkout_and_restore_preserve_process_input_recovery() { + use super::super::{InputEvent, handle_input}; + + for initially_overflowed in [false, true] { + let mut app = app(); + app.start_session("source".into()); + app.editor.insert_str("parked draft"); + if initially_overflowed { + handle_input(&mut app, InputEvent::Overflow); + handle_input(&mut app, InputEvent::RecoveryReady(true)); + } + prepare_branch(&mut app); + assert_eq!(app.input_overflow, initially_overflowed); + assert_eq!(app.input_recovery_ready, initially_overflowed); + assert_eq!(app.editor.text(), "original prompt"); + // These queue controls precede all modal dispatch, even while the + // source is parked. Restoring it must not restore old reader state. + handle_input(&mut app, InputEvent::Overflow); + handle_input(&mut app, InputEvent::RecoveryReady(true)); + app.abandon_branch(); + assert!(app.input_overflow); + assert!(app.input_recovery_ready); + assert_eq!(app.editor.text(), "parked draft"); + + prepare_branch(&mut app); + let received_at = Instant::now(); + handle_input(&mut app, InputEvent::Resumed(received_at)); + app.abandon_branch(); + assert!(!app.input_overflow); + assert!(!app.input_recovery_ready); + assert_eq!(app.last_key, Some(received_at)); + assert_eq!(app.editor.text(), "parked draft"); + } + } + + fn commit_branch_view(app: &mut App, child: &str) { + let arrival = Instant::now(); + app.last_key = Some(arrival - super::PASTE_GAP * 2); + let Action::SubmitPromptBranch { epoch, .. } = + app.handle_branch_key_at(press(KeyCode::Enter), arrival) + else { + panic!("submit") + }; + assert!(app.branch_submitted(epoch, child.into())); + } + + #[test] + fn branch_storage_state_is_visible_before_during_and_after_checkout() { + for initial in [(true, false), (false, true), (true, true)] { + for activate in [false, true] { + let mut app = app(); + app.start_session("source".into()); + app.activate_runtime_session(); + app.apply(Update::Runtime(RuntimeEvent::StorageStatus { + pending: initial.0, + exhausted: initial.1, + })); + prepare_branch(&mut app); + assert_eq!((app.storage_pending, app.storage_exhausted), initial); + // Even a child diagnostic-route change cannot hide global warnings. + app.apply(Update::Runtime(RuntimeEvent::SessionStarted { + session_id: "child".into(), + })); + for state in [(false, false), (true, true), (false, true), (true, false)] { + app.apply(Update::Runtime(RuntimeEvent::StorageStatus { + pending: state.0, + exhausted: state.1, + })); + assert_eq!((app.storage_pending, app.storage_exhausted), state); + let parked = &app.branch_draft.as_ref().unwrap().original; + assert_eq!((parked.storage_pending, parked.storage_exhausted), state); + } + if activate { + commit_branch_view(&mut app, "child"); + } else { + app.abandon_branch(); + } + assert!(app.storage_pending); + assert!(!app.storage_exhausted); + app.apply(Update::Runtime(RuntimeEvent::StorageStatus { + pending: false, + exhausted: false, + })); + assert!(!app.storage_pending); + assert!(!app.storage_exhausted); + } + } + } + + #[test] + fn branch_early_diagnostics_replay_after_activation_and_route_back_to_loaded_source() { + let mut app = app(); + let mut ingress = None; + app.start_session("source".into()); + app.activate_runtime_session(); + prepare_branch(&mut app); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::SessionStarted { + session_id: "child".into(), + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::CompactionStarted { + reason: "early".into(), + at: 0, + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + agent_event( + "child-agent", + "Child agent", + crate::events::SubagentStatus::Working, + None, + 1, + None, + (1, 1, None), + ), + ); + apply_stderr_runtime(&mut app, &mut ingress, child("call-1:compose:1", "shell")); + let parked = &app.branch_draft.as_ref().unwrap().original; + assert_eq!(parked.runtime_session_id.as_deref(), Some("source")); + assert!(!parked.compacting); + assert!(parked.agents.is_empty()); + assert!(!app.compacting); // provisional context is not the child yet + commit_branch_view(&mut app, "child"); + assert_eq!(app.runtime_session_id.as_deref(), Some("child")); + assert_eq!(app.branch_runtime_replay.as_ref().unwrap().len(), 3); + // A response is not the replay boundary: keep diagnostics even if ACP + // replay arrives later than stderr from the newly activated child. + apply_stderr_runtime( + &mut app, + &mut ingress, + child("call-1:compose:extra", "shell"), + ); + assert!(!app.compacting); + compose(&mut app, "shell({})"); + app.apply(Update::AvailableCommands { + session_id: "child".into(), + commands: Vec::new(), + }); + assert!(app.branch_runtime_replay.is_none()); + assert!(app.compacting); + assert!(app.agents.contains_key("child-agent")); + assert_eq!(app.call_mut("call-1").unwrap().children.len(), 2); + + // Loaded resume emits an ordered stderr source boundary before responding. + app.start_session("source".into()); + app.activate_runtime_session(); + compose(&mut app, "source shell({})"); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::SessionStarted { + session_id: "source".into(), + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::CompactionStarted { + reason: "source".into(), + at: 1, + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + agent_event( + "source-agent", + "Source agent", + crate::events::SubagentStatus::Working, + None, + 1, + None, + (1, 1, None), + ), + ); + apply_stderr_runtime(&mut app, &mut ingress, child("call-1:compose:2", "shell")); + assert_eq!(app.runtime_session_id.as_deref(), Some("source")); + assert!(app.compacting); + assert!(app.agents.contains_key("source-agent")); + assert!(!app.agents.contains_key("child-agent")); + assert_eq!(app.call_mut("call-1").unwrap().children.len(), 1); + } + + #[test] + fn branch_abandon_does_not_retarget_source_diagnostics_to_an_early_child() { + let mut app = app(); + let mut ingress = None; + app.start_session("source".into()); + app.activate_runtime_session(); + prepare_branch(&mut app); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::SessionStarted { + session_id: "child".into(), + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::CompactionStarted { + reason: "child".into(), + at: 0, + }, + ); + app.abandon_branch(); + assert_eq!(app.runtime_session_id.as_deref(), Some("source")); + assert!(!app.compacting); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::SessionStarted { + session_id: "source".into(), + }, + ); + apply_stderr_runtime( + &mut app, + &mut ingress, + RuntimeEvent::CompactionStarted { + reason: "source".into(), + at: 1, + }, + ); + assert!(app.compacting); + } + + #[test] + fn prompt_branch_abandon_restores_parked_editor_attachments_transcript_and_config() { + let mut app = app(); + app.start_session("source".into()); + app.push_block(Block::Agent("original answer".into())); + app.paste("unsent draft"); + app.attach( + PathBuf::from("/tmp/parked.png"), + "image/png", + AttachmentKind::Image, + 123, + ); + app.editor.move_left(); + let text = app.editor.text().to_string(); + let cursor = app.editor.cursor(); + let attachments = app.attachments.clone(); + let next_attachment = app.next_attachment; + app.reasoning_effort = "high".into(); + app.scroll = 7; + app.follow = false; + prepare_branch(&mut app); + assert!(app.editing_branch()); + assert!(app.blocks.is_empty()); + assert!(app.attachments.is_empty()); + assert_eq!(app.editor.text(), "original prompt"); + app.attach( + PathBuf::from("/tmp/rejected.png"), + "image/png", + AttachmentKind::Image, + 1, + ); + assert!(app.attachments.is_empty()); + app.paste(" edited"); + assert!(matches!(app.handle_key(press(KeyCode::Esc)), Action::None)); + assert!(!app.editing_branch()); + assert_eq!(app.session_id.as_deref(), Some("source")); + assert!(matches!(&app.blocks[0], Block::Agent(text) if text == "original answer")); + assert_eq!(app.editor.text(), text); + assert_eq!(app.editor.cursor(), cursor); + assert_eq!(app.attachments, attachments); + assert_eq!(app.next_attachment, next_attachment); + assert_eq!(app.reasoning_effort, "high"); + assert_eq!(app.scroll, 7); + assert!(!app.follow); + } + + #[test] + fn prompt_branch_stale_list_prepare_and_submit_cannot_replace_newer_draft_or_route() { + use crate::protocols::acp::prompt_branches::{PreparePromptBranchResponse, PromptBoundary}; + let mut app = app(); + app.start_session("source".into()); + let Action::ListPromptBranches { epoch } = app.open_branch_chooser() else { + panic!() + }; + app.abandon_branch(); + app.open_branch_chooser(); + app.branch_listed( + epoch, + Ok(vec![PromptBoundary { + address: "stale".into(), + text: "stale".into(), + historical: false, + }]), + ); + assert!(app.branch_chooser.as_ref().unwrap().boundaries.is_empty()); + app.branch_prepared( + epoch, + Ok(PreparePromptBranchResponse { + checkout_token: "stale".into(), + original_text: "stale".into(), + prefix: Vec::new(), + config_options: Vec::new(), + }), + ); + assert!(!app.editing_branch()); + app.abandon_branch(); + prepare_branch(&mut app); + let Action::SubmitPromptBranch { epoch, text, .. } = + app.handle_branch_key(press(KeyCode::Enter), false) + else { + panic!() + }; + assert_eq!(text, "original prompt"); + app.start_session("other".into()); + app.branch_submit_failed(epoch, "stale error".into()); + assert!(!app.branch_submitted(epoch, "stale-child".into())); + assert_eq!(app.session_id.as_deref(), Some("other")); + assert!(!app.editing_branch()); + } + + #[test] + fn prompt_branch_is_discoverable_without_stealing_navigator_search_characters() { + let mut app = app(); + app.start_session("source".into()); + app.paste("park this draft"); + app.open_navigation(); + assert!(matches!( + app.handle_navigation_key(press(KeyCode::Char('e')), false), + Action::None + )); + assert_eq!(app.navigation.dialog.as_ref().unwrap().query, "e"); + app.handle_navigation_key(press(KeyCode::Backspace), false); + app.paste("/branch"); + assert!(matches!( + app.handle_navigation_key(press(KeyCode::Enter), false), + Action::ListPromptBranches { .. } + )); + assert_eq!(app.editor.text(), "park this draft"); + app.abandon_branch(); + assert!(app.navigation.dialog.is_some()); + assert_eq!(app.editor.text(), "park this draft"); + } + + #[test] + fn prompt_branch_submit_is_text_only_and_never_an_ordinary_send() { + let mut app = app(); + app.start_session("source".into()); + prepare_branch(&mut app); + for code in [KeyCode::F(2), KeyCode::F(3)] { + assert!(matches!(app.handle_key(press(code)), Action::None)); + assert!(app.navigation.dialog.is_none()); + assert!(!app.queue_focused); + } + app.editor.clear(); + app.paste("/model literal edited text"); + assert!(matches!( + app.handle_branch_key(press(KeyCode::Enter), true), + Action::None + )); + assert!(app.editor.text().ends_with('\n')); + let Action::SubmitPromptBranch { + epoch, + checkout_token, + text, + } = app.handle_branch_key(press(KeyCode::Enter), false) + else { + panic!("must submit branch, not prompt/config") + }; + assert_eq!(checkout_token, "checkout-one"); + assert_eq!(text, "/model literal edited text\n"); + assert!(app.branch_submitting()); + app.paste("ignored while committing"); + app.handle_key(press(KeyCode::Esc)); + assert!(app.branch_submitting()); + assert_eq!(app.editor.text(), text); + app.branch_submit_failed(epoch, "try again".into()); + assert!(app.editing_branch()); + assert!(!app.branch_submitting()); + app.handle_key(press(KeyCode::Esc)); + assert_eq!(app.session_id.as_deref(), Some("source")); + } + #[test] fn model_switch_dialog_actions_preserve_input_and_selection() { use crate::protocols::acp::model_switch::{Decision, Warning}; diff --git a/src/tui/command.rs b/src/tui/command.rs index 0fe2440..5d4fcc9 100644 --- a/src/tui/command.rs +++ b/src/tui/command.rs @@ -15,6 +15,7 @@ enum Kind { Effort, Agents, Transcript, + Branch, Login, } @@ -89,6 +90,11 @@ const LOCAL_COMMANDS: &[Spec] = &[ description: "Browse the transcript (F3)", kind: Kind::Transcript, }, + Spec { + token: "/branch", + description: "Edit a previous text prompt in a new session", + kind: Kind::Branch, + }, Spec { token: "/login", description: "Authenticate with the agent", @@ -106,6 +112,7 @@ pub enum Parsed<'a> { Effort { value: Option<&'a str> }, Agents, Transcript, + Branch, Login { method_id: Option<&'a str> }, Prompt(&'a str), } @@ -142,7 +149,8 @@ pub fn parse(input: &str, login_available: bool) -> Parsed<'_> { Kind::Effort => Parsed::Effort { value: prompt }, Kind::Agents if prompt.is_none() => Parsed::Agents, Kind::Transcript if prompt.is_none() => Parsed::Transcript, - Kind::Agents | Kind::Transcript => Parsed::Prompt(input), + Kind::Branch if prompt.is_none() => Parsed::Branch, + Kind::Agents | Kind::Transcript | Kind::Branch => Parsed::Prompt(input), Kind::Login => Parsed::Login { method_id: prompt.map(str::trim), }, @@ -302,6 +310,15 @@ mod tests { ); } + #[test] + fn branch_is_an_exact_local_command_with_completion() { + assert_eq!(parse("/branch"), Parsed::Branch); + assert_eq!(parse("/branch "), Parsed::Branch); + assert_eq!(parse("/branch extra"), Parsed::Prompt("/branch extra")); + assert_eq!(known_token("/branch", &[]), Some(0..7)); + assert_eq!(completions("/br", 3, &[])[0].name, "/branch"); + } + #[test] fn transcript_is_an_exact_highlighted_local_command_without_arguments() { for login_available in [false, true] { @@ -405,6 +422,7 @@ mod tests { "/effort", "/agents", "/transcript", + "/branch", "/compact", ] ); diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 59a954b..6b27215 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -65,6 +65,11 @@ use crate::{ tools::mcp::CredentialStorage, }; +use crate::protocols::acp::prompt_branches::{ + ListPromptBranchesRequest, ListPromptBranchesResponse, PreparePromptBranchRequest, + PreparePromptBranchResponse, SubmitPromptBranchRequest, SubmitPromptBranchResponse, +}; + use app::{ Action, App, Attachment, AttachmentKind, EffortChoice, ModelChoice, SubmittedPrompt, Update, UserImage, @@ -171,6 +176,42 @@ struct ActiveSessionRoute { generation: u64, } +enum BranchResponse { + Listed(Result), + Prepared(Result), + Submitted(Result), +} + +struct BranchCompletion { + generation: u64, + epoch: u64, + response: BranchResponse, +} + +// A child ID is unknown until submit responds. Buffer its replay and early live +// notifications under the same lock used to install the route at completion. +#[derive(Default)] +struct BranchReplayBuffer { + request: Option<(u64, u64)>, + notifications: Vec, +} + +impl BranchCompletion { + fn is_current(&self, route: &ActiveSessionRoute, app: &App) -> bool { + self.generation == route.generation && self.epoch == app.branch_epoch + } +} + +impl BranchReplayBuffer { + fn finish(&mut self, generation: u64, epoch: u64) -> Vec { + if self.request != Some((generation, epoch)) { + return Vec::new(); + } + self.request = None; + std::mem::take(&mut self.notifications) + } +} + struct QueuedUpdate { generation: Option, update: Update, @@ -192,6 +233,26 @@ impl QueuedUpdate { } } +// Stderr and ACP are independent streams. Capture identity while reading stderr, +// not when applying its queued events after a possibly newer ACP activation. +// Markers are stream boundaries, never commands to switch the visible session. +fn runtime_diagnostic_update( + route: &mut Option, + event: events::RuntimeEvent, +) -> Option { + match event { + events::RuntimeEvent::SessionStarted { session_id } => { + *route = Some(session_id); + None + } + event @ events::RuntimeEvent::StorageStatus { .. } => Some(Update::Runtime(event)), + event => route.as_ref().map(|session_id| Update::RoutedRuntime { + session_id: session_id.clone(), + event, + }), + } +} + fn transition_route(route: &Arc>, id: String) { if let Ok(mut route) = route.lock() { route.id = id; @@ -224,7 +285,9 @@ fn apply_pending_updates( let Some(update) = accept_queued_update(route, queued) else { continue; }; - if let Update::ConfigOptions(options) = &update { + if let Update::ConfigOptions(options) = &update + && !app.editing_branch() + { refresh_config_state(app, Some(options)); } app.apply(update); @@ -863,6 +926,9 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let stderr = child.stderr.take().ok_or("could not open Kit stderr")?; let transport = ByteStreams::new(stdin.compat_write(), stdout.compat()); let (updates_tx, mut updates_rx) = mpsc::unbounded_channel(); + let (branch_tx, mut branch_rx) = mpsc::unbounded_channel::(); + let branch_replay = Arc::new(Mutex::new(BranchReplayBuffer::default())); + let notification_branch_replay = Arc::clone(&branch_replay); // The agent's own diagnostics are the only explanation of a failed start, // so they are kept aside as well as shown in the log pane. @@ -871,9 +937,16 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let diagnostics = updates_tx.clone(); let stderr_task = tokio::spawn(async move { let mut lines = BufReader::new(stderr).lines(); + let mut runtime_route = None; while let Ok(Some(line)) = lines.next_line().await { let update = match events::parse(&line) { - Some(event) => Update::Runtime(event), + Some(event) => { + let Some(update) = runtime_diagnostic_update(&mut runtime_route, event) + else { + continue; + }; + update + } None if line.starts_with("A2A listening on ") => { Update::A2aAddress(line.trim_start_matches("A2A listening on ").to_string()) } @@ -938,7 +1011,14 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( .v2() .on_receive_notification( async move |notification: UpdateSessionNotification, _cx| { + let mut replay = notification_branch_replay.lock().expect("branch replay lock"); let current = notification_session.lock().ok().map(|route| route.clone()); + if replay.request.is_some() + && current.as_ref().is_some_and(|route| route.id != notification.session_id.to_string()) + { + replay.notifications.push(notification); + return Ok(()); + } for update in current.as_ref().map_or_else(Vec::new, |route| { translate_for_session(notification, &route.id) }) { @@ -1224,6 +1304,8 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( active.id = active_session_id.clone(); } app.start_session(active_session_id.clone()); + app.activate_runtime_session(); + let mut preserved_branch_sources = std::collections::HashSet::new(); let storage_shutdown = crate::resilient_fs::shutdown_token(); let result: Result<(), agent_client_protocol::Error> = async { loop { @@ -1247,6 +1329,44 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( }; match action { Action::Quit => return Ok(()), + Action::ListPromptBranches { epoch } => { + let generation = transition_session.lock().expect("session route lock").generation; + let connection = connection.clone(); + let session_id = session_id.clone(); + let sender = branch_tx.clone(); + tokio::spawn(async move { + let result = connection.send_request(ListPromptBranchesRequest { session_id }) + .block_task().await.map_err(|error| error.message.to_string()); + let _ = sender.send(BranchCompletion { generation, epoch, response: BranchResponse::Listed(result) }); + }); + } + Action::PreparePromptBranch { epoch, address } => { + let generation = transition_session.lock().expect("session route lock").generation; + let connection = connection.clone(); + let session_id = session_id.clone(); + let sender = branch_tx.clone(); + tokio::spawn(async move { + let result = connection.send_request(PreparePromptBranchRequest { session_id, address }) + .block_task().await.map_err(|error| error.message.to_string()); + let _ = sender.send(BranchCompletion { generation, epoch, response: BranchResponse::Prepared(result) }); + }); + } + Action::SubmitPromptBranch { epoch, checkout_token, text } => { + let generation = transition_session.lock().expect("session route lock").generation; + { + let mut replay = branch_replay.lock().expect("branch replay lock"); + replay.request = Some((generation, epoch)); + replay.notifications.clear(); + } + let connection = connection.clone(); + let session_id = session_id.clone(); + let sender = branch_tx.clone(); + tokio::spawn(async move { + let result = connection.send_request(SubmitPromptBranchRequest { session_id, checkout_token, text }) + .block_task().await.map_err(|error| error.message.to_string()); + let _ = sender.send(BranchCompletion { generation, epoch, response: BranchResponse::Submitted(result) }); + }); + } Action::Submit { prompt, inject } => { let blocks = match prompt_blocks(&prompt) { Ok(blocks) => blocks, @@ -1361,6 +1481,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( transition_route(&transition_session, persisted_id.clone()); images.clear(); app.start_session(persisted_id); + app.activate_runtime_session(); refresh_config_state(&mut app, Some(&config_options)); if let Some(prompt) = first_prompt { let outcome = connection @@ -1458,7 +1579,11 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( // lock that the OS proves no live Kit process still holds. if let Err(error) = crate::session::load(&root, &requested_id) .and_then(|_| { - crate::session::remove_stale_lock(&root, &requested_id) + if preserved_branch_sources.contains(&requested_id) { + Ok(()) + } else { + crate::session::remove_stale_lock(&root, &requested_id) + } }) { app.note(format!("could not resume session: {error}")); @@ -1481,10 +1606,10 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( images.clear(); app.start_session(requested_id.clone()); match request_resume(&connection, session_id.clone(), root.clone()).await { - Ok(response) => refresh_config_state( - &mut app, - Some(&response.config_options), - ), + Ok(response) => { + app.activate_runtime_session(); + refresh_config_state(&mut app, Some(&response.config_options)); + }, Err(error) => { session_id = previous_session_id; transition_route( @@ -1501,6 +1626,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( .await; match restored { Ok(response) => { + app.activate_runtime_session(); refresh_config_state( &mut app, Some(&response.config_options), @@ -1764,6 +1890,43 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( } } }, + Some(completion) = branch_rx.recv() => { + let mut replay = branch_replay.lock().expect("branch replay lock"); + let current = completion.is_current(&transition_session.lock().expect("session route lock"), &app); + let submitted = matches!(completion.response, BranchResponse::Submitted(_)); + let buffered = if submitted { + replay.finish(completion.generation, completion.epoch) + } else { Vec::new() }; + if !current { continue; } + match completion.response { + BranchResponse::Listed(result) => app.branch_listed(completion.epoch, result.map(|response| response.boundaries)), + BranchResponse::Prepared(result) => { + app.branch_prepared(completion.epoch, result); + images.clear(); + } + BranchResponse::Submitted(Err(error)) => app.branch_submit_failed(completion.epoch, error), + BranchResponse::Submitted(Ok(response)) => { + let child_id = response.session_id.to_string(); + if app.branch_submitted(completion.epoch, child_id.clone()) { + preserved_branch_sources.insert(session_id.to_string()); + session_id = response.session_id; + transition_route(&transition_session, child_id.clone()); + images.clear(); + refresh_config_state(&mut app, Some(&response.config_options)); + // Replay precedes activation, just as with resume. No ordinary + // prompt request: submit already persisted the edited prompt. + for notification in buffered { + for update in translate_for_session(notification, &child_id) { + if let Update::ConfigOptions(options) = &update { + refresh_config_state(&mut app, Some(options)); + } + app.apply(update); + } + } + } + } + } + }, update = updates_rx.recv() => match update { Some(update) => apply_pending_updates( &mut app, @@ -2073,7 +2236,10 @@ fn handle(app: &mut App, received: ReceivedEvent) -> Action { if app.model_switch.is_some() { return Action::None; } - if app.navigation.dialog.is_some() { + if app.navigation.dialog.is_some() + || app.branch_chooser.is_some() + || app.editing_branch() + { app.paste(&text); return Action::None; } @@ -3353,6 +3519,121 @@ mod tests { ); } + #[tokio::test] + async fn delayed_stderr_cannot_retarget_loaded_source_or_misattribute_child_events() { + use crate::events::{EVENT_MARKER, RuntimeEvent}; + use crate::tui::app::Block; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, Lines}; + + async fn ingest( + lines: &mut Lines>, + ingress: &mut Option, + tx: &tokio::sync::mpsc::UnboundedSender, + ) { + for _ in 0..4 { + let line = lines.next_line().await.unwrap().unwrap(); + if let Some(update) = + super::runtime_diagnostic_update(ingress, crate::events::parse(&line).unwrap()) + { + tx.send(QueuedUpdate::global(update)).unwrap(); + } + } + } + + // Exercise both delays: stderr already read but queued behind ACP resume, + // and child bytes still in the pipe when ACP resume activates the source. + for read_before_resume in [false, true] { + let root = tempfile::tempdir().unwrap(); + let mut app = App::new( + root.path().into(), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.start_session("child".into()); + app.activate_runtime_session(); + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "child".into(), + generation: 0, + })); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let (mut writer, reader) = tokio::io::duplex(4096); + let mut lines = BufReader::new(reader).lines(); + let mut ingress = None; + for (session, call, pending) in [ + ("child", "compose:child", true), + ("source", "compose:source", false), + ] { + for event in [ + RuntimeEvent::SessionStarted { + session_id: session.into(), + }, + RuntimeEvent::CompactionStarted { + reason: session.into(), + at: 0, + }, + RuntimeEvent::ChildStarted { + call: call.into(), + tool: "shell".into(), + summary: session.into(), + at: 0, + }, + RuntimeEvent::StorageStatus { + pending, + exhausted: pending, + }, + ] { + writer + .write_all( + format!("{EVENT_MARKER}{}\n", serde_json::to_string(&event).unwrap()) + .as_bytes(), + ) + .await + .unwrap(); + } + if session == "child" { + if read_before_resume { + ingest(&mut lines, &mut ingress, &tx).await; + } + transition_route(&route, "source".into()); + app.start_session("source".into()); + app.activate_runtime_session(); + app.apply(Update::ToolStarted { + id: "compose".into(), + title: "compose".into(), + kind: wire::ToolKind::Other, + script: Some("shell({})".into()), + backgrounded: false, + }); + if !read_before_resume { + ingest(&mut lines, &mut ingress, &tx).await; + } + } else { + // Loaded resume must emit this source marker on stderr even + // though it reuses an actor and does not start a new driver. + ingest(&mut lines, &mut ingress, &tx).await; + } + let first = rx.try_recv().unwrap(); + apply_pending_updates(&mut app, &route, &mut rx, first); + assert_eq!(app.session_id.as_deref(), Some("source")); + assert_eq!(app.compacting, session == "source"); + assert_eq!( + (app.storage_pending, app.storage_exhausted), + (pending, pending) + ); + let Block::Tool(compose) = &app.blocks[0] else { + panic!("compose") + }; + if session == "child" { + assert!(compose.children.is_empty()); + } else { + assert_eq!(compose.children.len(), 1); + assert_eq!(compose.children[0].call, "compose:source"); + } + } + } + } + #[test] fn queued_updates_are_applied_in_bounded_bursts() { let root = tempfile::tempdir().unwrap(); @@ -3469,6 +3750,70 @@ mod tests { )); } + #[test] + fn prompt_branch_completions_require_both_route_generation_and_draft_epoch() { + let mut app = App::new( + PathBuf::from("/tmp"), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.branch_epoch = 3; + let route = super::ActiveSessionRoute { + id: "source".into(), + generation: 7, + }; + for response in [ + super::BranchResponse::Listed(Err("list".into())), + super::BranchResponse::Prepared(Err("prepare".into())), + super::BranchResponse::Submitted(Err("submit".into())), + ] { + let mut completion = super::BranchCompletion { + generation: 7, + epoch: 3, + response, + }; + assert!(completion.is_current(&route, &app)); + completion.epoch = 2; + assert!(!completion.is_current(&route, &app)); + completion.epoch = 3; + completion.generation = 6; + assert!(!completion.is_current(&route, &app)); + } + } + + #[test] + fn prompt_branch_replay_waits_for_child_route_and_stale_submit_does_not_drain_new_buffer() { + let mut buffer = super::BranchReplayBuffer { + request: Some((7, 3)), + notifications: vec![UpdateSessionNotification::new( + "child", + SessionUpdate::UserMessage( + UserMessage::new("edited-prompt") + .content(vec![ContentBlock::Text(TextContent::new("edited"))]), + ), + )], + }; + assert!(buffer.finish(7, 2).is_empty()); + assert_eq!(buffer.notifications.len(), 1); + let route = Arc::new(Mutex::new(super::ActiveSessionRoute { + id: "source".into(), + generation: 7, + })); + let notifications = buffer.finish(7, 3); + assert!(buffer.request.is_none()); + transition_route(&route, "child".into()); + let active = route.lock().unwrap(); + let updates: Vec<_> = notifications + .into_iter() + .flat_map(|notification| translate_for_session(notification, &active.id)) + .collect(); + assert!( + matches!(updates.as_slice(), [Update::UserMessage { id, text, append: false, .. }] if id == "edited-prompt" && text == "edited") + ); + assert_eq!(active.generation, 8); + } + #[test] fn translates_accepted_user_messages_and_foreground_state() { let user = UpdateSessionNotification::new( diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 28809b8..cfa1238 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -73,6 +73,7 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { && app.blocks.is_empty() && app.pending_steers.is_empty() && !app.editing_steer() + && !app.editing_branch() && !app.queue_focused && !app.show_logs; @@ -93,8 +94,14 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { let minimum_rows = 1 + 3 + logs_rows + pending_rows + prompt_rows + 1; let rainbow_fits = frame.area().height >= minimum_rows.saturating_add(1); let header_rows = 1 + u16::from(!app.blocks.is_empty() && rainbow_fits); - let [header, body, logs, pending, prompt, status] = Layout::vertical([ + let warning_rows = if app.editing_branch() { + (super::app::BRANCH_WARNING.len() as u16 / frame.area().width.max(1) + 2).min(7) + } else { + 0 + }; + let [header, warning, body, logs, pending, prompt, status] = Layout::vertical([ Constraint::Length(header_rows), + Constraint::Length(warning_rows), Constraint::Min(3), Constraint::Length(logs_rows), Constraint::Length(pending_rows), @@ -104,13 +111,23 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { .areas(frame.area()); draw_header(frame, app, header); + if app.editing_branch() { + frame.render_widget( + Paragraph::new(super::app::BRANCH_WARNING) + .style(Style::default().fg(theme::warn_color())) + .wrap(ratatui::widgets::Wrap { trim: false }), + warning, + ); + } draw_body(frame, app, images, body); if app.show_logs { draw_logs(frame, app, logs); } draw_pending_steers(frame, app, pending); let viewport = draw_prompt(frame, app, prompt); - draw_command_popup(frame, app, prompt); + if !app.editing_branch() { + draw_command_popup(frame, app, prompt); + } draw_status(frame, app, status); (prompt, viewport, false) }; @@ -128,6 +145,28 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { if app.navigation.dialog.is_some() { draw_navigation(frame, app, &navigation_matches); } + if app.branch_chooser.is_some() { + draw_branch_chooser(frame, app); + } + // Durability stays visible on the start screen and over session pickers. + // Pending data belongs to the process, not the currently selected session. + if app.storage_pending || app.storage_exhausted { + let area = frame.area(); + let warning = if app.storage_exhausted { + " Storage exhausted: shutting down; unpersisted data is at risk" + } else { + " Memory-only storage: awaiting disk recovery; data at risk on exit" + }; + frame.render_widget( + Paragraph::new(warning).style(Style::default().fg(theme::warn_color())), + Rect::new( + area.x, + area.bottom().saturating_sub(1), + area.width, + u16::from(area.height > 0), + ), + ); + } // This warning must not expire or disappear behind a modal while the reader // discards input. Esc acknowledgement is handled by the reader, not the UI. if app.input_overflow { @@ -150,25 +189,6 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { warning, ); } - // Durability stays visible on the start screen and over session pickers. - // Pending data belongs to the process, not the currently selected session. - if app.storage_pending || app.storage_exhausted { - let area = frame.area(); - let warning = if app.storage_exhausted { - " Storage exhausted: shutting down; unpersisted data is at risk" - } else { - " Memory-only storage: awaiting disk recovery; data at risk on exit" - }; - frame.render_widget( - Paragraph::new(warning).style(Style::default().fg(theme::warn_color())), - Rect::new( - area.x, - area.bottom().saturating_sub(1), - area.width, - u16::from(area.height > 0), - ), - ); - } } #[derive(Clone, Copy)] @@ -661,6 +681,86 @@ fn draw_effort_dialog(frame: &mut Frame<'_>, app: &App) { ); } +/// Only backend-approved text boundaries are selectable, including archived ones. +fn draw_branch_chooser(frame: &mut Frame<'_>, app: &App) { + let chooser = app.branch_chooser.as_ref().expect("checked above"); + let outer = frame.area(); + let width = outer.width.saturating_sub(4).clamp(1, 96); + let height = outer.height.saturating_sub(2).clamp(1, 24); + let area = Rect::new( + outer.x + outer.width.saturating_sub(width) / 2, + outer.y + outer.height.saturating_sub(height) / 2, + width, + height, + ); + let panel = Panel::bordered().title(" prompt checkout · text only · source preserved "); + let inner = panel.inner(area); + frame.render_widget(Clear, area); + frame.render_widget(panel, area); + let [warning, entries, help] = Layout::vertical([ + Constraint::Length( + (super::app::BRANCH_WARNING.len() as u16 / inner.width.max(1) + 2).min(7), + ), + Constraint::Min(1), + Constraint::Length(2), + ]) + .areas(inner); + frame.render_widget( + Paragraph::new(super::app::BRANCH_WARNING) + .style(Style::default().fg(theme::warn_color())) + .wrap(ratatui::widgets::Wrap { trim: false }), + warning, + ); + let rows = if chooser.pending { + vec![Line::from("Loading checkout… Esc cancels")] + } else if chooser.boundaries.is_empty() { + vec![Line::from("No eligible text prompts in this session.")] + } else { + let start = chooser + .selected + .saturating_sub(entries.height.saturating_sub(1) as usize); + chooser + .boundaries + .iter() + .enumerate() + .skip(start) + .take(entries.height as usize) + .map(|(index, boundary)| { + let preview: String = boundary + .text + .chars() + .map(|c| if c.is_control() { ' ' } else { c }) + .take(inner.width as usize) + .collect(); + Line::styled( + format!( + "{} {}{} {}", + if index == chooser.selected { + "›" + } else { + " " + }, + if boundary.historical { + "[archived] " + } else { + "" + }, + boundary.address, + preview + ), + if index == chooser.selected { + theme::bold(theme::accent_color()) + } else { + theme::dim() + }, + ) + }) + .collect() + }; + frame.render_widget(Paragraph::new(rows), entries); + frame.render_widget(Paragraph::new("↑/↓ select · Enter edit in provisional draft · Esc back\nNo branch is created until the edited draft is submitted."), help); +} + /// A read-only index of display text. Only visible entries build previews. fn draw_navigation(frame: &mut Frame<'_>, app: &App, matches: &[usize]) { let outer = frame.area(); @@ -759,7 +859,7 @@ fn draw_navigation(frame: &mut Frame<'_>, app: &App, matches: &[usize]) { frame.render_widget( Paragraph::new(vec![ Line::from(Span::styled( - "↑/↓ select · Enter reveal · Esc/F3 close", + "↑/↓ · Enter reveal · /branch ↵ edit prompt · Esc/F3 close", theme::dim(), )), Line::from(Span::styled( @@ -2252,7 +2352,11 @@ fn draw_prompt(frame: &mut Frame<'_>, app: &App, area: Rect) -> PromptViewport { theme::faint() }; let block = Panel::bordered() - .title(if app.editing_steer() { + .title(if app.branch_submitting() { + " creating prompt branch… " + } else if app.editing_branch() { + " provisional prompt checkout · Enter create branch · Esc abandon " + } else if app.editing_steer() { " editing pending · Enter save · Esc cancel " } else { "" @@ -2444,7 +2548,11 @@ fn draw_status(frame: &mut Frame<'_>, app: &App, area: Rect) { Style::default().fg(theme::warn_color()), )); } - let hints = if app.editing_steer() { + let hints = if app.branch_submitting() { + "creating branch — source preserved " + } else if app.editing_branch() { + "⏎ create branch esc abandon ⇧⏎ newline · text only " + } else if app.editing_steer() { "⏎ save edit esc restore draft " } else if app.queue_focused && !app.pending_steers.is_empty() { if app.can_replace_steer @@ -3112,6 +3220,93 @@ mod tests { app.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); } + #[test] + fn prompt_branch_chooser_and_provisional_editor_show_warning_and_archived_label() { + use crate::protocols::acp::prompt_branches::{PreparePromptBranchResponse, PromptBoundary}; + let mut app = navigation_app(Vec::new()); + app.session_id = Some("source".into()); + app.branch_chooser = Some(crate::tui::app::BranchChooser { + boundaries: vec![PromptBoundary { + address: "opaque-boundary".into(), + text: "editable text".into(), + historical: true, + }], + selected: 0, + pending: false, + }); + let output = render(&mut app, 100, 24); + assert!(output.contains("[archived] opaque-boundary"), "{output}"); + assert!( + output.contains("Only conversation context changes."), + "{output}" + ); + assert!(output.contains("not rolled back."), "{output}"); + app.branch_prepared( + app.branch_epoch, + Ok(PreparePromptBranchResponse { + checkout_token: "checkout".into(), + original_text: "editable text".into(), + prefix: Vec::new(), + config_options: Vec::new(), + }), + ); + let output = render(&mut app, 100, 24); + assert!(output.contains("provisional prompt checkout"), "{output}"); + assert!(output.contains("editable text"), "{output}"); + assert!( + output.contains("Only conversation context changes."), + "{output}" + ); + assert!(output.contains("not rolled back."), "{output}"); + assert!(output.contains("Esc abandon"), "{output}"); + } + + #[test] + fn input_overflow_warning_is_last_overlay_in_branch_views() { + use crate::protocols::acp::prompt_branches::PreparePromptBranchResponse; + use crate::tui::{handle_input, input::InputEvent}; + + let mut app = navigation_app(Vec::new()); + app.session_id = Some("source".into()); + app.branch_chooser = Some(crate::tui::app::BranchChooser { + boundaries: Vec::new(), + selected: 0, + pending: false, + }); + app.storage_pending = true; + handle_input(&mut app, InputEvent::Overflow); + for draft in [false, true] { + if draft { + app.branch_prepared( + app.branch_epoch, + Ok(PreparePromptBranchResponse { + checkout_token: "checkout".into(), + original_text: "editable text".into(), + prefix: Vec::new(), + config_options: Vec::new(), + }), + ); + } + for ready in [false, true] { + handle_input(&mut app, InputEvent::RecoveryReady(ready)); + let output = render(&mut app, 60, 14); + assert!( + output.contains("Input overflow: input discarded"), + "{output}" + ); + assert!(output.contains("Check draft before sending"), "{output}"); + assert!( + output.contains(if ready { + "Press Esc to resume" + } else { + "Waiting for quiet input" + }), + "{output}" + ); + } + } + } + #[test] fn navigation_reveal_uses_wrapped_prefix_and_reanchors_on_resize() { let mut app = navigation_app(vec![