diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index 582e6612..5c2028d0 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -40,6 +40,13 @@ impl CompletionState { pub const MAX_VISIBLE: usize = 6; } +struct RollbackSubmitState { + text: String, + segments: Vec, + block_start: usize, + turn_before: usize, +} + pub struct App { pub pod_name: String, pub connected: bool, @@ -91,6 +98,12 @@ pub struct App { /// Top entry index of the task pane's visible window. Clamped on /// render so it never points past the end of the list. pub task_pane_scroll: usize, + /// Local submit state kept until the accepted run either completes + /// normally or reports that the empty assistant turn was rolled back. + pending_submit_rollback: Option, + /// Last rolled-back submit that could not be restored because the + /// composer already contained unsent user input. + last_rolled_back_input: Option>, } impl App { @@ -120,6 +133,8 @@ impl App { task_store: TaskStore::new(), task_pane_open: false, task_pane_scroll: 0, + pending_submit_rollback: None, + last_rolled_back_input: None, } } @@ -339,7 +354,15 @@ impl App { // TurnHeader / UserMessage blocks are pushed in response to // `Event::UserMessage` (single source of truth, shared by every // client subscribed to the Pod). Locally we only clear the - // input buffer and forward the method. + // input buffer and forward the method, while remembering enough + // local state to undo the visible submit if the Pod reports that + // the accepted run produced no assistant output and was rolled back. + self.pending_submit_rollback = Some(RollbackSubmitState { + text: Segment::flatten_to_text(&segments), + segments: segments.clone(), + block_start: self.blocks.len(), + turn_before: self.turn_index, + }); self.input.clear(); Some(Method::Run { input: segments }) } @@ -670,22 +693,22 @@ impl App { self.push_error(format!("[{code:?}] {message}")); } Event::RunEnd { result } => { - self.blocks.push(Block::TurnStats { - requests: self.run_requests, - upload_tokens: self.run_upload_tokens, - output_tokens: self.run_output_tokens, - }); - self.set_pod_status(match result { - RunResult::Paused => PodStatus::Paused, - RunResult::Finished | RunResult::LimitReached | RunResult::RolledBack => { - PodStatus::Idle - } - }); - self.run_requests = 0; - self.run_upload_tokens = 0; - self.run_output_tokens = 0; - self.current_tool = None; - self.assistant_streaming = false; + if matches!(result, RunResult::RolledBack) { + self.handle_rolled_back_run(); + } else { + self.blocks.push(Block::TurnStats { + requests: self.run_requests, + upload_tokens: self.run_upload_tokens, + output_tokens: self.run_output_tokens, + }); + self.pending_submit_rollback = None; + self.reset_run_state(match result { + RunResult::Paused => PodStatus::Paused, + RunResult::Finished | RunResult::LimitReached | RunResult::RolledBack => { + PodStatus::Idle + } + }); + } } Event::CompactStart => { self.blocks.push(Block::Compact(CompactEvent::Streaming { @@ -770,6 +793,44 @@ impl App { } } + fn reset_run_state(&mut self, status: PodStatus) { + self.set_pod_status(status); + self.run_requests = 0; + self.run_upload_tokens = 0; + self.run_output_tokens = 0; + self.current_tool = None; + self.assistant_streaming = false; + } + + fn handle_rolled_back_run(&mut self) { + let hint = if let Some(state) = self.pending_submit_rollback.take() { + self.blocks + .truncate(state.block_start.min(self.blocks.len())); + self.turn_index = state.turn_before; + if self.input.is_empty() { + self.input.replace_with_segments(&state.segments); + self.completion = None; + self.last_rolled_back_input = None; + "Rolled back empty assistant turn; restored your input.".to_owned() + } else { + let preview = rollback_input_preview(&state.text); + self.last_rolled_back_input = Some(state.segments); + format!( + "Rolled back empty assistant turn; composer was not empty, kept submitted input in backup: {preview}" + ) + } + } else { + "Rolled back empty assistant turn; no local submitted input was available to restore." + .to_owned() + }; + self.reset_run_state(PodStatus::Idle); + self.blocks.push(Block::Alert { + level: AlertLevel::Warn, + source: AlertSource::Pod, + message: hint, + }); + } + fn append_assistant_text(&mut self, text: &str) { if self.assistant_streaming { if let Some(Block::AssistantText { text: existing }) = self.blocks.last_mut() { @@ -1098,6 +1159,16 @@ fn strip_cat_n_prefix(formatted: &str) -> String { out } +fn rollback_input_preview(text: &str) -> String { + const MAX_CHARS: usize = 80; + let mut one_line = text.replace('\n', "⏎"); + if one_line.chars().count() > MAX_CHARS { + one_line = one_line.chars().take(MAX_CHARS).collect::(); + one_line.push('…'); + } + one_line +} + /// True if the submitted segment list carries no user-visible content /// (only whitespace / newlines, no paste, no typed atoms). Used to /// decide whether an empty Enter should be a no-op or trigger a @@ -1439,6 +1510,133 @@ mod completion_flow_tests { assert!(app.completion.as_ref().unwrap().entries.is_empty()); } + #[test] + fn rolled_back_run_restores_input_and_removes_submit_blocks() { + let mut app = App::new("test".into()); + let submitted = submit_text(&mut app, "please wait"); + assert_eq!(input_text(&app), ""); + + app.handle_pod_event(Event::UserMessage { + segments: submitted, + }); + // Simulate run-derived attachment display after the submitted user line. + app.blocks.push(Block::SystemMessage { + text: "[File: README.md]".into(), + }); + app.handle_pod_event(Event::TurnStart { turn: 1 }); + app.handle_pod_event(Event::Usage { + input_tokens: Some(100), + output_tokens: Some(0), + cache_read_input_tokens: Some(40), + }); + app.handle_pod_event(Event::RunEnd { + result: RunResult::RolledBack, + }); + + assert_eq!(input_text(&app), "please wait"); + assert_eq!(app.turn_index, 0); + assert!(app.blocks.iter().all(|b| !matches!( + b, + Block::TurnHeader { .. } + | Block::UserMessage { .. } + | Block::SystemMessage { .. } + | Block::TurnStats { .. } + ))); + assert!(warning_contains(&app, "restored your input")); + assert!(matches!(app.pod_status, PodStatus::Idle)); + assert!(!app.running); + assert!(!app.paused); + assert_eq!(app.run_requests, 0); + assert_eq!(app.run_upload_tokens, 0); + assert_eq!(app.run_output_tokens, 0); + assert!(app.current_tool.is_none()); + } + + #[test] + fn rolled_back_run_does_not_overwrite_existing_unsent_input() { + let mut app = App::new("test".into()); + let submitted = submit_text(&mut app, "original submit"); + app.handle_pod_event(Event::UserMessage { + segments: submitted, + }); + for c in "draft while running".chars() { + app.insert_char(c); + } + + app.handle_pod_event(Event::RunEnd { + result: RunResult::RolledBack, + }); + + assert_eq!(input_text(&app), "draft while running"); + assert_eq!( + Segment::flatten_to_text(app.last_rolled_back_input.as_ref().unwrap()), + "original submit" + ); + assert!(warning_contains(&app, "composer was not empty")); + assert!(app.blocks.iter().all(|b| !matches!( + b, + Block::TurnHeader { .. } | Block::UserMessage { .. } | Block::TurnStats { .. } + ))); + } + + #[test] + fn non_rolled_back_run_end_keeps_submitted_blocks_and_does_not_restore_input() { + for result in [RunResult::Paused, RunResult::Finished] { + let mut app = App::new("test".into()); + let submitted = submit_text(&mut app, "normal run"); + app.handle_pod_event(Event::UserMessage { + segments: submitted, + }); + app.handle_pod_event(Event::RunEnd { result }); + + assert_eq!(input_text(&app), ""); + assert!( + app.blocks + .iter() + .any(|b| matches!(b, Block::TurnHeader { .. })) + ); + assert!( + app.blocks + .iter() + .any(|b| matches!(b, Block::UserMessage { .. })) + ); + assert!( + app.blocks + .iter() + .any(|b| matches!(b, Block::TurnStats { .. })) + ); + assert!(!warning_contains(&app, "Rolled back empty assistant turn")); + assert!(app.last_rolled_back_input.is_none()); + } + } + + fn submit_text(app: &mut App, text: &str) -> Vec { + for c in text.chars() { + app.insert_char(c); + } + match app.submit_input() { + Some(Method::Run { input }) => input, + other => panic!("expected Run, got {other:?}"), + } + } + + fn input_text(app: &App) -> String { + Segment::flatten_to_text(&app.input.submit_segments()) + } + + fn warning_contains(app: &App, needle: &str) -> bool { + app.blocks.iter().any(|block| { + matches!( + block, + Block::Alert { + level: AlertLevel::Warn, + message, + .. + } if message.contains(needle) + ) + }) + } + #[test] fn snapshot_renders_system_message_block_from_session_start() { let mut app = App::new("test".into()); diff --git a/crates/tui/src/input.rs b/crates/tui/src/input.rs index 5f7951c7..50e58546 100644 --- a/crates/tui/src/input.rs +++ b/crates/tui/src/input.rs @@ -168,6 +168,56 @@ impl InputBuffer { self.cursor = 0; } + pub fn is_empty(&self) -> bool { + self.atoms.is_empty() + } + + /// Replace the whole composer with protocol segments previously emitted + /// by [`submit_segments`](Self::submit_segments), preserving typed chips + /// and placing the cursor at the end of the restored input. + pub fn replace_with_segments(&mut self, segments: &[protocol::Segment]) { + self.atoms.clear(); + for segment in segments { + match segment { + protocol::Segment::Text { content } => { + self.atoms.extend(content.chars().map(Atom::Char)); + } + protocol::Segment::Paste { + id, + chars, + lines, + content, + } => { + self.next_paste_id = self.next_paste_id.max(id.saturating_add(1).max(1)); + self.atoms.push(Atom::Paste(PasteRef { + id: *id, + chars: *chars as usize, + lines: *lines as usize, + content: content.clone(), + })); + } + protocol::Segment::FileRef { path } => { + self.atoms + .push(Atom::FileRef(FileRefAtom { path: path.clone() })); + } + protocol::Segment::KnowledgeRef { slug } => { + self.atoms + .push(Atom::KnowledgeRef(KnowledgeRefAtom { slug: slug.clone() })); + } + protocol::Segment::WorkflowInvoke { slug } => { + self.atoms.push(Atom::WorkflowInvoke(WorkflowInvokeAtom { + slug: slug.clone(), + })); + } + protocol::Segment::Unknown => { + self.atoms + .extend("[unknown input segment]".chars().map(Atom::Char)); + } + } + } + self.cursor = self.atoms.len(); + } + pub fn insert_char(&mut self, c: char) { self.atoms.insert(self.cursor, Atom::Char(c)); self.cursor += 1;