Fix resetting Idle state after canceling agent

Author: stippi

crates/code_assistant/src/session/instance.rs

@@ -429,19 +429,22 @@ impl UserInterface for ProxyUI {
                 // Update activity state to waiting for response
                 self.update_activity_state(SessionActivityState::WaitingForResponse);
             }
-            UiEvent::StreamingStopped { .. } => {
+            UiEvent::StreamingStopped { cancelled, .. } => {
                 // Clear fragment buffer when LLM request ends - fragments are now part of message history
                 if let Ok(mut buffer) = self.fragment_buffer.lock() {
                     buffer.clear();
                 }
-                // Update activity state back to agent running (it will be set to idle when agent completes)
-                let current_state = self.session_activity_state.lock().unwrap().clone();
-                if matches!(
-                    current_state,
-                    SessionActivityState::WaitingForResponse
-                        | SessionActivityState::RateLimited { .. }
-                ) {
-                    self.update_activity_state(SessionActivityState::AgentRunning);
+                // Only update activity state back to agent running if streaming was not cancelled
+                // If cancelled, the agent task will end soon and set state to Idle
+                if !cancelled {
+                    let current_state = self.session_activity_state.lock().unwrap().clone();
+                    if matches!(
+                        current_state,
+                        SessionActivityState::WaitingForResponse
+                            | SessionActivityState::RateLimited { .. }
+                    ) {
+                        self.update_activity_state(SessionActivityState::AgentRunning);
+                    }
                 }
             }
             _ => {}

crates/code_assistant/src/session/manager.rs

@@ -280,19 +280,27 @@ impl SessionManager {
             debug!("Starting agent for session {}", session_id_clone);
             let result = agent.run_single_iteration().await;
 
-            // Set session state back to Idle when agent completes and broadcast
+            // Always set session state back to Idle when agent task ends (regardless of success/failure/cancellation)
+            debug!("Agent task ending for session {}, setting state to Idle", session_id_clone);
             if let Ok(mut state) = activity_state_ref.lock() {
                 *state = crate::session::instance::SessionActivityState::Idle;
             }
 
-            let _ = ui_clone
+            // Always broadcast the state change to UI
+            let send_result = ui_clone
                 .send_event(crate::ui::UiEvent::UpdateSessionActivityState {
                     session_id: session_id_clone.clone(),
                     activity_state: crate::session::instance::SessionActivityState::Idle,
                 })
                 .await;
 
-            debug!("Agent completed for session {}", session_id_clone);
+            if let Err(e) = send_result {
+                debug!("Failed to send UpdateSessionActivityState event for session {}: {}", session_id_clone, e);
+            } else {
+                debug!("Successfully sent UpdateSessionActivityState(Idle) event for session {}", session_id_clone);
+            }
+
+            debug!("Agent completed for session {} with result: {:?}", session_id_clone, result.is_ok());
             result
         });