Skip to content

Commit 60b331e

Browse files
vanzuevanzue
authored
Fix turn cancellation lifecycle (#832)
* Fix turn cancellation lifecycle Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Address cancellation review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Document cancellation quiescence boundary Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Clarify cancellation draft preservation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Close cancellation channel races Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Settle renamed pre-dispatch cancellations Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Simplify turn cancellation ownership Replace tab-keyed cancellation bookkeeping with prompt-scoped tokens, preserve producer quiescence across transport and session lifecycle changes, and harden exact session routing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec * Make repeated stop cancellation idempotent Treat the cancellation-drain state as stoppable without emitting a duplicate cancellation notice or a misleading nothing-to-stop message. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec --------- Copilot-Session: 8d6e172c-4c82-4d13-b97a-f557c7e980ec
1 parent c1a25cd commit 60b331e

16 files changed

Lines changed: 3834 additions & 827 deletions

‎tools/wta/Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tools/wta/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ path = "src/main.rs"
1212
[dependencies]
1313
agent-client-protocol = { version = "1.3.0" }
1414
tokio = { version = "1", features = ["full"] }
15-
tokio-util = { version = "0.7", features = ["compat"] }
15+
tokio-util = { version = "0.7", features = ["compat", "rt"] }
1616
async-trait = "0.1"
1717
anyhow = "1"
1818
serde_json = "1"

‎tools/wta/src/app.rs‎

Lines changed: 164 additions & 36 deletions
Large diffs are not rendered by default.

‎tools/wta/src/app/autofix.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,11 @@ impl App {
289289
// lookup so a tab with no ACP session yet still gets the prompt
290290
// queued correctly (the ACP layer creates the session lazily when
291291
// it processes the prompt).
292-
self.turn_submit_prompt_for_tab(&target_tab_id, submitted);
292+
self.turn_submit_prompt_for_tab_with_cancellation(
293+
&target_tab_id,
294+
submitted,
295+
prompt.cancellation_token(),
296+
);
293297
tracing::info!(target: "autofix", pane_id = %notification.pane_id, tab_id = %target_tab_id, generation = new_gen, "sending auto-fix prompt");
294298
let _ = self.prompt_tx.send(prompt);
295299

@@ -529,7 +533,7 @@ impl App {
529533
let (turn_matches, state_matches) = autofix_pane_matches(tab, pane_id);
530534

531535
if turn_matches {
532-
self.turn_cancel_for_tab(&target_tab_id);
536+
self.request_turn_cancel_for_tab(&target_tab_id);
533537
} else if state_matches {
534538
let tab = self.tab_mut(&target_tab_id);
535539
tab.autofix.generation = tab.autofix.generation.wrapping_add(1);

‎tools/wta/src/app/tab_state.rs‎

Lines changed: 98 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,13 @@ pub struct UserInputState {
241241
Option<tokio::sync::oneshot::Sender<crate::agent_tools::user_input::UserInputResponse>>,
242242
}
243243

244+
pub(crate) struct ActivePromptCancellation {
245+
pub prompt_id: u64,
246+
pub token: tokio_util::sync::CancellationToken,
247+
pub session_id: Option<String>,
248+
pub attachment_valid: bool,
249+
}
250+
244251
impl UserInputState {
245252
pub fn selection_count(&self) -> usize {
246253
self.request.choices.len() + usize::from(self.request.allow_freeform)
@@ -585,6 +592,7 @@ pub struct TabSession {
585592
// Explicit per-turn lifecycle. Source of truth in the new state machine
586593
// (see `doc/specs/turn-state-refactor.md`).
587594
pub turn: TurnState,
595+
pub(crate) active_prompt_cancellation: Option<ActivePromptCancellation>,
588596
pub activity_frame: usize,
589597
/// Ephemeral ACP thought text shown only until visible assistant output
590598
/// or another structured activity begins. It never enters `messages` or
@@ -684,6 +692,74 @@ impl TabSession {
684692
)?
685693
}
686694

695+
pub(crate) fn set_prompt_cancellation(
696+
&mut self,
697+
prompt_id: u64,
698+
token: tokio_util::sync::CancellationToken,
699+
) {
700+
self.active_prompt_cancellation = Some(ActivePromptCancellation {
701+
prompt_id,
702+
token,
703+
session_id: self.session_id.clone(),
704+
attachment_valid: true,
705+
});
706+
}
707+
708+
pub(crate) fn can_attach_prompt_session(&self, prompt_id: u64, session_id: &str) -> bool {
709+
self.active_prompt_cancellation
710+
.as_ref()
711+
.is_some_and(|active| {
712+
active.prompt_id == prompt_id
713+
&& active.attachment_valid
714+
&& active
715+
.session_id
716+
.as_deref()
717+
.is_none_or(|known| known == session_id)
718+
})
719+
}
720+
721+
pub(crate) fn bind_active_prompt_session(&mut self, prompt_id: u64, session_id: &str) {
722+
if let Some(active) = self.active_prompt_cancellation.as_mut().filter(|active| {
723+
active.prompt_id == prompt_id && active.attachment_valid && active.session_id.is_none()
724+
}) {
725+
active.session_id = Some(session_id.to_string());
726+
}
727+
}
728+
729+
pub(crate) fn invalidate_active_prompt_attachment(&mut self) {
730+
if let Some(active) = self.active_prompt_cancellation.as_mut() {
731+
active.attachment_valid = false;
732+
}
733+
}
734+
735+
pub(crate) fn cancel_active_prompt(&self, prompt_id: u64) {
736+
if let Some(active) = self
737+
.active_prompt_cancellation
738+
.as_ref()
739+
.filter(|active| active.prompt_id == prompt_id)
740+
{
741+
active.token.cancel();
742+
}
743+
}
744+
745+
pub(crate) fn active_prompt_matches_session(&self, prompt_id: u64, session_id: &str) -> bool {
746+
self.active_prompt_cancellation
747+
.as_ref()
748+
.is_some_and(|active| {
749+
active.prompt_id == prompt_id && active.session_id.as_deref() == Some(session_id)
750+
})
751+
}
752+
753+
pub(crate) fn finish_active_prompt(&mut self, prompt_id: u64) {
754+
if self
755+
.active_prompt_cancellation
756+
.as_ref()
757+
.is_some_and(|active| active.prompt_id == prompt_id)
758+
{
759+
self.active_prompt_cancellation = None;
760+
}
761+
}
762+
687763
pub(crate) fn cached_completed_turn_height(
688764
&self,
689765
index: usize,
@@ -935,6 +1011,25 @@ impl TabSession {
9351011
}
9361012

9371013
pub fn clear_chat_history(&mut self) {
1014+
let cancellation_barrier = match &self.turn {
1015+
TurnState::Submitted(prompt)
1016+
| TurnState::Streaming { prompt }
1017+
| TurnState::Surfaced {
1018+
prompt,
1019+
end_pending: true,
1020+
..
1021+
} => Some(prompt.id),
1022+
TurnState::Cancelling { prompt_id } => Some(*prompt_id),
1023+
TurnState::Idle
1024+
| TurnState::Surfaced {
1025+
end_pending: false, ..
1026+
} => None,
1027+
};
1028+
if let Some(prompt_id) = cancellation_barrier {
1029+
self.cancel_active_prompt(prompt_id);
1030+
} else {
1031+
self.active_prompt_cancellation = None;
1032+
}
9381033
self.messages.clear();
9391034
self.clear_streaming_thought();
9401035
self.permission.clear();
@@ -947,7 +1042,9 @@ impl TabSession {
9471042
self.timing_note = None;
9481043
self.selection_visible_pending = false;
9491044
self.clear_completed_turn_selection();
950-
self.turn = TurnState::Idle;
1045+
self.turn = cancellation_barrier
1046+
.map(|prompt_id| TurnState::Cancelling { prompt_id })
1047+
.unwrap_or(TurnState::Idle);
9511048
self.clear_recommendations();
9521049
self.attachments
9531050
.remove_tokens_from_input(&mut self.input, &mut self.cursor_pos);

‎tools/wta/src/app/turn_state.rs‎

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,11 @@ pub enum TurnState {
2929
outcome: TurnOutcome,
3030
end_pending: bool,
3131
},
32+
/// The user requested cancellation, but the ACP prompt has not reached
33+
/// its terminal boundary yet. New prompts remain blocked and all
34+
/// turn-scoped updates are discarded until the correlated cancellation
35+
/// settlement arrives.
36+
Cancelling { prompt_id: u64 },
3237
}
3338

3439
impl Default for TurnState {
@@ -100,7 +105,9 @@ impl TurnState {
100105
match self {
101106
TurnState::Idle => true,
102107
TurnState::Surfaced { end_pending, .. } => !*end_pending,
103-
_ => false,
108+
TurnState::Submitted(_)
109+
| TurnState::Streaming { .. }
110+
| TurnState::Cancelling { .. } => false,
104111
}
105112
}
106113

@@ -114,7 +121,25 @@ impl TurnState {
114121
TurnState::Surfaced {
115122
end_pending: true, ..
116123
} => true,
117-
_ => false,
124+
TurnState::Idle
125+
| TurnState::Surfaced {
126+
end_pending: false, ..
127+
}
128+
| TurnState::Cancelling { .. } => false,
129+
}
130+
}
131+
132+
pub fn is_cancelling(&self) -> bool {
133+
matches!(self, TurnState::Cancelling { .. })
134+
}
135+
136+
pub fn prompt_id(&self) -> Option<u64> {
137+
match self {
138+
TurnState::Idle => None,
139+
TurnState::Submitted(prompt)
140+
| TurnState::Streaming { prompt }
141+
| TurnState::Surfaced { prompt, .. } => Some(prompt.id),
142+
TurnState::Cancelling { prompt_id } => Some(*prompt_id),
118143
}
119144
}
120145

@@ -125,7 +150,7 @@ impl TurnState {
125150
/// request itself proves the Agent is still waiting; only `Idle` means
126151
/// there is no turn left to service.
127152
pub fn can_service_agent_request(&self) -> bool {
128-
!matches!(self, TurnState::Idle)
153+
!matches!(self, TurnState::Idle | TurnState::Cancelling { .. })
129154
}
130155

131156
/// The surfaced recommendation set, if the outcome is a card.
@@ -142,7 +167,7 @@ impl TurnState {
142167
/// Prompt info for the in-flight or just-surfaced turn.
143168
pub fn prompt(&self) -> Option<&SubmittedPrompt> {
144169
match self {
145-
TurnState::Idle => None,
170+
TurnState::Idle | TurnState::Cancelling { .. } => None,
146171
TurnState::Submitted(p) => Some(p),
147172
TurnState::Streaming { prompt } => Some(prompt),
148173
TurnState::Surfaced { prompt, .. } => Some(prompt),
@@ -154,7 +179,7 @@ impl TurnState {
154179
/// `App::apply_prompt_target_resolved`).
155180
pub fn prompt_mut(&mut self) -> Option<&mut SubmittedPrompt> {
156181
match self {
157-
TurnState::Idle => None,
182+
TurnState::Idle | TurnState::Cancelling { .. } => None,
158183
TurnState::Submitted(p) => Some(p),
159184
TurnState::Streaming { prompt } => Some(prompt),
160185
TurnState::Surfaced { prompt, .. } => Some(prompt),
@@ -291,6 +316,17 @@ mod tests {
291316
assert!(s.recommendations().is_some());
292317
}
293318

319+
#[test]
320+
fn cancelling_blocks_new_prompts_and_drops_turn_updates() {
321+
let s = TurnState::Cancelling { prompt_id: 1 };
322+
assert!(!s.is_idle());
323+
assert!(!s.accepts_new_prompt());
324+
assert!(!s.is_in_flight());
325+
assert!(s.is_cancelling());
326+
assert!(!s.can_service_agent_request());
327+
assert!(s.prompt().is_none());
328+
}
329+
294330
#[test]
295331
fn surfaced_chat_has_no_recommendation() {
296332
let s = TurnState::Surfaced {

‎tools/wta/src/app_contracts/event.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ pub enum AppEvent {
3636
SessionAttached {
3737
tab_id: String,
3838
session_id: String,
39+
/// Present only for a session lazily created by this exact prompt.
40+
/// Lifecycle-created sessions (`/new`, `session/load`, startup
41+
/// fallback) are not prompt-owned.
42+
prompt_id: Option<u64>,
3943
available_models: Vec<AcpModelInfo>,
4044
current_model_id: Option<String>,
4145
},
@@ -92,6 +96,11 @@ pub enum AppEvent {
9296
tab_id: String,
9397
message: String,
9498
},
99+
PromptError {
100+
tab_id: String,
101+
prompt_id: u64,
102+
message: String,
103+
},
95104
TabSystemMessage {
96105
tab_id: String,
97106
message: String,
@@ -122,6 +131,11 @@ pub enum AppEvent {
122131
/// The helper's pipe to wta-master closed. A retained helper reconnects
123132
/// its existing immutable binding over the stable pipe.
124133
MasterDisconnected,
134+
/// The ACP client has conclusively retired its transport and no prompt
135+
/// or lifecycle task owned by that client can produce more events.
136+
/// Releases cancellation barriers for both dispatched and still-queued
137+
/// prompts.
138+
AgentTransportRetired,
125139
AgentSoftStop {
126140
session_id: String,
127141
reason: crate::protocol::acp::soft_stop::SoftStopReason,
@@ -151,6 +165,10 @@ pub enum AppEvent {
151165
AgentMessageEnd {
152166
session_id: String,
153167
},
168+
PromptCancellationSettled {
169+
prompt_id: u64,
170+
started: bool,
171+
},
154172
TimingMetric {
155173
session_id: String,
156174
note: String,

0 commit comments

Comments
 (0)