Add cancellation flow for support sessions

Server changes:
- Allow cancelling connected codes (not just pending)
- Reject agent connections with cancelled codes
- Periodic cancellation check during active sessions
- Send Disconnect message when code is cancelled

Agent changes:
- Detect cancellation via Disconnect message
- Show Windows MessageBox to notify user
- Exit cleanly without reconnecting for support sessions

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
2025-12-28 15:30:43 -07:00
parent f408667a3f
commit 8246d135f9
4 changed files with 150 additions and 8 deletions

View File

@@ -73,6 +73,28 @@ async fn handle_agent_connection(
) {
info!("Agent connected: {} ({})", agent_name, agent_id);
let (mut ws_sender, mut ws_receiver) = socket.split();
// If a support code was provided, check if it's valid
if let Some(ref code) = support_code {
// Check if the code is cancelled or invalid
if support_codes.is_cancelled(code).await {
warn!("Agent tried to connect with cancelled code: {}", code);
// Send disconnect message to agent
let disconnect_msg = proto::Message {
payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
reason: "Support session was cancelled by technician".to_string(),
})),
};
let mut buf = Vec::new();
if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
let _ = ws_sender.send(Message::Binary(buf.into())).await;
}
let _ = ws_sender.close().await;
return;
}
}
// Register the agent and get channels
let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;
@@ -85,12 +107,16 @@ async fn handle_agent_connection(
support_codes.link_session(code, session_id).await;
}
let (mut ws_sender, mut ws_receiver) = socket.split();
// Use Arc<Mutex> for sender so we can use it from multiple places
let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
let ws_sender_input = ws_sender.clone();
let ws_sender_cancel = ws_sender.clone();
// Task to forward input events from viewers to agent
let input_forward = tokio::spawn(async move {
while let Some(input_data) = input_rx.recv().await {
if ws_sender.send(Message::Binary(input_data.into())).await.is_err() {
let mut sender = ws_sender_input.lock().await;
if sender.send(Message::Binary(input_data.into())).await.is_err() {
break;
}
}
@@ -99,6 +125,34 @@ async fn handle_agent_connection(
let sessions_cleanup = sessions.clone();
let support_codes_cleanup = support_codes.clone();
let support_code_cleanup = support_code.clone();
let support_code_check = support_code.clone();
let support_codes_check = support_codes.clone();
// Task to check for cancellation every 2 seconds
let cancel_check = tokio::spawn(async move {
let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
loop {
interval.tick().await;
if let Some(ref code) = support_code_check {
if support_codes_check.is_cancelled(code).await {
info!("Support code {} was cancelled, disconnecting agent", code);
// Send disconnect message
let disconnect_msg = proto::Message {
payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
reason: "Support session was cancelled by technician".to_string(),
})),
};
let mut buf = Vec::new();
if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
let mut sender = ws_sender_cancel.lock().await;
let _ = sender.send(Message::Binary(buf.into())).await;
let _ = sender.close().await;
}
break;
}
}
}
});
// Main loop: receive frames from agent and broadcast to viewers
while let Some(msg) = ws_receiver.next().await {
@@ -135,12 +189,15 @@ async fn handle_agent_connection(
// Cleanup
input_forward.abort();
cancel_check.abort();
sessions_cleanup.remove_session(session_id).await;
// Mark support code as completed if one was used
// Mark support code as completed if one was used (unless cancelled)
if let Some(ref code) = support_code_cleanup {
support_codes_cleanup.mark_completed(code).await;
info!("Support code {} marked as completed", code);
if !support_codes_cleanup.is_cancelled(code).await {
support_codes_cleanup.mark_completed(code).await;
info!("Support code {} marked as completed", code);
}
}
info!("Session {} ended", session_id);

View File

@@ -176,11 +176,11 @@ impl SupportCodeManager {
}
}
/// Cancel a code
/// Cancel a code (works for both pending and connected)
pub async fn cancel_code(&self, code: &str) -> bool {
let mut codes = self.codes.write().await;
if let Some(support_code) = codes.get_mut(code) {
if support_code.status == CodeStatus::Pending {
if support_code.status == CodeStatus::Pending || support_code.status == CodeStatus::Connected {
support_code.status = CodeStatus::Cancelled;
return true;
}
@@ -188,6 +188,18 @@ impl SupportCodeManager {
false
}
/// Check if a code is cancelled
pub async fn is_cancelled(&self, code: &str) -> bool {
let codes = self.codes.read().await;
codes.get(code).map(|c| c.status == CodeStatus::Cancelled).unwrap_or(false)
}
/// Check if a code is valid for connection (exists and is pending)
pub async fn is_valid_for_connection(&self, code: &str) -> bool {
let codes = self.codes.read().await;
codes.get(code).map(|c| c.status == CodeStatus::Pending).unwrap_or(false)
}
/// List all codes (for dashboard)
pub async fn list_codes(&self) -> Vec<SupportCode> {
let codes = self.codes.read().await;