fix: distinguish agent completion from startup and session changes (#4457)

refs #3338
refs #3990
refs #3916

1 files changed, 30 insertions, 0 deletions

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<crate::protocol::ClientShellSnapshot>),
     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<crate::protocol::endpoint::EndpointAgentViewProjection, _> =
             serde_json::from_str(data)
@@ -78,6 +84,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 {
             source: "example.views".into(),