From 28360107c1bfab57fb0de99d74bfa48e74423ca8 Mon Sep 17 00:00:00 2001 From: Can Celik Date: Mon, 21 Sep 2026 20:47:14 +0300 Subject: fix: distinguish agent completion from startup and session changes (#4457) refs #3338 refs #3990 refs #3916 --- src/client/endpoint/control.rs | 30 ++++ src/client/mod.rs | 10 ++ src/client/shell/endpoint_agent_state.rs | 235 +++++++++++++++++++++++++++++-- src/client/shell/endpoints.rs | 19 ++- 4 files changed, 282 insertions(+), 12 deletions(-) (limited to 'src/client') diff --git a/src/client/endpoint/control.rs b/src/client/endpoint/control.rs index 6eef8dc8..f2fe1b76 100644 --- a/src/client/endpoint/control.rs +++ b/src/client/endpoint/control.rs @@ -9,6 +9,7 @@ pub(crate) struct DecodedAgentViewProjection { pub(crate) enum EndpointControlMessage { HealthPong, AgentViewProjection(DecodedAgentViewProjection), + AgentCompletions(crate::protocol::endpoint::EndpointAgentCompletions), Snapshot(Box), Ignored, } @@ -20,6 +21,11 @@ pub(crate) fn decode_endpoint_control( if kind == crate::protocol::endpoint::HEALTH_PONG_KIND { return Ok(EndpointControlMessage::HealthPong); } + if kind == crate::protocol::endpoint::AGENT_COMPLETIONS_KIND { + return Ok(serde_json::from_str(data) + .map(EndpointControlMessage::AgentCompletions) + .unwrap_or(EndpointControlMessage::Ignored)); + } if kind == crate::protocol::endpoint::AGENT_VIEW_PROJECTION_KIND { let Ok(projection): Result = serde_json::from_str(data) @@ -77,6 +83,30 @@ mod tests { )); } + #[test] + fn completion_guard_optional_control_round_trips_without_changing_snapshot_codec() { + let projection = crate::protocol::endpoint::EndpointAgentCompletions { + boot_id: "boot".into(), + revision: 3, + completions: [("pane".into(), 7)].into_iter().collect(), + }; + let crate::protocol::ServerMessage::EndpointControl { kind, data } = + crate::protocol::endpoint::agent_completions_message(&projection).unwrap() + else { + panic!("expected optional control"); + }; + let EndpointControlMessage::AgentCompletions(decoded) = + decode_endpoint_control(&kind, &data).unwrap() + else { + panic!("expected completion projection"); + }; + assert_eq!(decoded, projection); + assert!(matches!( + decode_endpoint_control(&kind, "invalid").unwrap(), + EndpointControlMessage::Ignored + )); + } + #[test] fn agent_view_projection_decodes_and_validates_the_view() { let view = crate::api::schema::AgentViewSetParams { diff --git a/src/client/mod.rs b/src/client/mod.rs index b9335e06..68e569f9 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1913,6 +1913,16 @@ async fn run_client_loop( } continue; } + Ok(endpoint::EndpointControlMessage::AgentCompletions(projection)) => { + if let Some(shell) = state.shell.as_mut() { + shell.set_endpoint_agent_completions( + &endpoint_id, + generation, + projection, + ); + } + continue; + } Ok(endpoint::EndpointControlMessage::Ignored) => { debug!(%kind, "ignoring unknown endpoint control message"); continue; diff --git a/src/client/shell/endpoint_agent_state.rs b/src/client/shell/endpoint_agent_state.rs index f8ba7a9b..1ebcd35d 100644 --- a/src/client/shell/endpoint_agent_state.rs +++ b/src/client/shell/endpoint_agent_state.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use crate::api::schema::AgentStatus; use crate::protocol::{ClientShellAgent, ClientShellSnapshot, PaneSurfaceFrame}; @@ -7,13 +7,49 @@ use crate::protocol::{ClientShellAgent, ClientShellSnapshot, PaneSurfaceFrame}; pub(super) struct EndpointAgentPresentation { boot_id: Option, acknowledged: HashMap, + completed: HashMap, + working: HashSet, + pending_completions: Option<( + Option, + crate::protocol::endpoint::EndpointAgentCompletions, + )>, } impl EndpointAgentPresentation { + pub(super) fn receive_completions( + &mut self, + generation: Option, + completions: crate::protocol::endpoint::EndpointAgentCompletions, + ) { + if self + .pending_completions + .as_ref() + .is_some_and(|(current_generation, current)| { + *current_generation == generation + && current.boot_id == completions.boot_id + && current.revision >= completions.revision + }) + { + return; + } + self.pending_completions = Some((generation, completions)); + } + + #[cfg(test)] pub(super) fn project_snapshot(&mut self, snapshot: &mut ClientShellSnapshot) { + self.project_snapshot_for_generation(snapshot, None); + } + + pub(super) fn project_snapshot_for_generation( + &mut self, + snapshot: &mut ClientShellSnapshot, + generation: Option, + ) { if self.boot_id.as_deref() != Some(snapshot.boot_id.as_str()) { self.boot_id = Some(snapshot.boot_id.clone()); self.acknowledged.clear(); + self.completed.clear(); + self.working.clear(); self.acknowledged.extend( snapshot .agents @@ -21,13 +57,59 @@ impl EndpointAgentPresentation { .map(|agent| (agent.pane_id.clone(), agent.state_change_seq)), ); } - self.acknowledged.retain(|pane_id, _| { - snapshot - .agents - .iter() - .any(|agent| &agent.pane_id == pane_id) - }); + let pane_ids: HashSet<&str> = snapshot + .agents + .iter() + .map(|agent| agent.pane_id.as_str()) + .collect(); + self.acknowledged + .retain(|pane_id, _| pane_ids.contains(pane_id.as_str())); + self.completed + .retain(|pane_id, _| pane_ids.contains(pane_id.as_str())); + self.working + .retain(|pane_id| pane_ids.contains(pane_id.as_str())); + let completions = self + .pending_completions + .take() + .filter(|(received_generation, projection)| { + *received_generation == generation + && projection.boot_id == snapshot.boot_id + && projection.revision == snapshot.revision + }) + .map(|(_, projection)| projection.completions); for agent in &mut snapshot.agents { + match agent.agent_status { + AgentStatus::Working => { + self.working.insert(agent.pane_id.clone()); + self.completed.remove(&agent.pane_id); + } + AgentStatus::Blocked => { + self.completed.remove(&agent.pane_id); + } + AgentStatus::Idle | AgentStatus::Done => { + let observed_work = self.working.remove(&agent.pane_id); + let completed = completions.as_ref().map_or_else( + || { + observed_work + || self.completed.get(&agent.pane_id) + == Some(&agent.state_change_seq) + }, + |completions| { + completions.get(&agent.pane_id) == Some(&agent.state_change_seq) + }, + ); + if completed { + self.completed + .insert(agent.pane_id.clone(), agent.state_change_seq); + } else { + self.completed.remove(&agent.pane_id); + } + } + _ => { + self.working.remove(&agent.pane_id); + self.completed.remove(&agent.pane_id); + } + } agent.agent_status = self.projected_status(agent); } project_aggregate_status(snapshot); @@ -72,9 +154,11 @@ impl EndpointAgentPresentation { } pub(super) fn seen(&self, agent: &ClientShellAgent) -> bool { - self.acknowledged - .get(&agent.pane_id) - .is_some_and(|sequence| *sequence >= agent.state_change_seq) + self.completed.get(&agent.pane_id).is_none_or(|completion| { + self.acknowledged + .get(&agent.pane_id) + .is_some_and(|sequence| sequence >= completion) + }) } fn projected_status(&self, agent: &ClientShellAgent) -> AgentStatus { @@ -119,7 +203,9 @@ fn project_aggregate_status(snapshot: &mut ClientShellSnapshot) { #[cfg(test)] mod tests { use super::*; - use crate::protocol::{FrameData, PaneSurfacePane, SurfaceRect}; + use crate::protocol::{ + endpoint::EndpointAgentCompletions, FrameData, PaneSurfacePane, SurfaceRect, + }; fn agent(status: AgentStatus, sequence: u64) -> ClientShellAgent { ClientShellAgent { @@ -197,6 +283,133 @@ mod tests { assert_eq!(snapshot.agents[0].agent_status, AgentStatus::Idle); } + fn assert_idle_sequence(states: &[(AgentStatus, u64)]) { + let mut presentation = EndpointAgentPresentation::default(); + let mut projected = AgentStatus::Unknown; + for (revision, &(status, seq)) in states.iter().enumerate() { + let mut snapshot = snapshot(status, seq, revision as u64 + 1); + presentation.project_snapshot(&mut snapshot); + projected = snapshot.agents[0].agent_status; + } + assert_eq!(projected, AgentStatus::Idle, "{states:?}"); + } + + fn completions(boot: &str, revision: u64, seq: Option) -> EndpointAgentCompletions { + EndpointAgentCompletions { + boot_id: boot.into(), + revision, + completions: seq + .map(|seq| ("agent-pane".into(), seq)) + .into_iter() + .collect(), + } + } + + #[test] + fn completion_guard_initial_unknown_to_idle_does_not_project_done() { + assert_idle_sequence(&[(AgentStatus::Unknown, 0), (AgentStatus::Idle, 1)]); + } + + #[test] + fn completion_guard_new_idle_agent_is_not_completed_work() { + let mut presentation = EndpointAgentPresentation::default(); + let mut baseline = snapshot(AgentStatus::Idle, 0, 1); + baseline.agents.clear(); + presentation.project_snapshot(&mut baseline); + let mut ready = snapshot(AgentStatus::Idle, 1, 2); + + presentation.project_snapshot(&mut ready); + + assert_eq!(ready.agents[0].agent_status, AgentStatus::Idle); + } + + #[test] + fn completion_guard_session_rebind_does_not_project_done() { + assert_idle_sequence(&[ + (AgentStatus::Idle, 4), + (AgentStatus::Unknown, 5), + (AgentStatus::Idle, 6), + ]); + } + + #[test] + fn completion_guard_startup_blocker_does_not_project_done() { + assert_idle_sequence(&[ + (AgentStatus::Unknown, 0), + (AgentStatus::Blocked, 1), + (AgentStatus::Idle, 2), + ]); + } + + #[test] + fn completion_guard_distinguishes_coalesced_work_from_rebind() { + for completion in [None, Some(6)] { + let mut presentation = EndpointAgentPresentation::default(); + let mut baseline = snapshot(AgentStatus::Idle, 4, 1); + presentation.project_snapshot(&mut baseline); + presentation.receive_completions(None, completions(&baseline.boot_id, 2, completion)); + let mut settled = snapshot(AgentStatus::Idle, 6, 2); + + presentation.project_snapshot(&mut settled); + + let expected = if completion.is_some() { + AgentStatus::Done + } else { + AgentStatus::Idle + }; + assert_eq!(settled.agents[0].agent_status, expected); + presentation.project_snapshot(&mut settled); + assert_eq!(settled.agents[0].agent_status, expected); + } + } + + #[test] + fn completion_guard_server_suppression_overrides_client_observed_work() { + let mut presentation = EndpointAgentPresentation::default(); + let mut baseline = snapshot(AgentStatus::Working, 1, 1); + presentation.project_snapshot(&mut baseline); + presentation.receive_completions(None, completions(&baseline.boot_id, 2, None)); + let mut settled = snapshot(AgentStatus::Idle, 2, 2); + + presentation.project_snapshot(&mut settled); + + assert_eq!(settled.agents[0].agent_status, AgentStatus::Idle); + } + + #[test] + fn completion_guard_companion_is_scoped_to_boot_revision_and_connection() { + for (boot, revision, generation, seq) in [ + ("old-boot", 2, Some(1), 6), + ("endpoint-boot", 1, Some(1), 6), + ("endpoint-boot", 2, Some(0), 6), + ("endpoint-boot", 2, Some(1), 5), + ] { + let mut presentation = EndpointAgentPresentation::default(); + let mut baseline = snapshot(AgentStatus::Idle, 4, 1); + presentation.project_snapshot_for_generation(&mut baseline, Some(1)); + presentation.receive_completions(generation, completions(boot, revision, Some(seq))); + let mut settled = snapshot(AgentStatus::Idle, 6, 2); + + presentation.project_snapshot_for_generation(&mut settled, Some(1)); + + assert_eq!(settled.agents[0].agent_status, AgentStatus::Idle); + } + } + + #[test] + fn completion_guard_new_boot_baselines_existing_completions() { + let mut presentation = EndpointAgentPresentation::default(); + let mut old = snapshot(AgentStatus::Working, 1, 1); + old.boot_id = "old-boot".into(); + presentation.project_snapshot(&mut old); + presentation.receive_completions(None, completions("endpoint-boot", 1, Some(6))); + let mut restored = snapshot(AgentStatus::Idle, 6, 1); + + presentation.project_snapshot(&mut restored); + + assert_eq!(restored.agents[0].agent_status, AgentStatus::Idle); + } + #[test] fn unpresented_working_completion_projects_done() { let mut presentation = EndpointAgentPresentation::default(); diff --git a/src/client/shell/endpoints.rs b/src/client/shell/endpoints.rs index 513842ba..dd6fa3cb 100644 --- a/src/client/shell/endpoints.rs +++ b/src/client/shell/endpoints.rs @@ -309,6 +309,23 @@ impl ClientShellState { .map(|snapshot| (snapshot.boot_id.as_str(), snapshot.revision)) } + pub(crate) fn set_endpoint_agent_completions( + &mut self, + endpoint_id: &ClientEndpointId, + generation: u64, + projection: crate::protocol::endpoint::EndpointAgentCompletions, + ) { + if let Some(endpoint) = self + .endpoints + .iter_mut() + .find(|endpoint| &endpoint.endpoint_id == endpoint_id) + { + endpoint + .agent_presentation + .receive_completions(Some(generation), projection); + } + } + pub(crate) fn set_endpoint_agent_view_projection_for_generation( &mut self, endpoint_id: &ClientEndpointId, @@ -546,7 +563,7 @@ impl ClientShellState { } self.endpoints[index] .agent_presentation - .project_snapshot(&mut snapshot); + .project_snapshot_for_generation(&mut snapshot, generation); let presented_surface = if acknowledge_surface && endpoint_id == &self.active_endpoint_id { self.pane_surface.as_ref() } else { -- cgit v1.2.3