Fix child pane reopening and overload retries

app/src/pane_group/child_agent/restoration.rs

@@ -107,13 +107,19 @@ impl PaneGroup {
         child_conversation_id: AIConversationId,
         ctx: &mut ViewContext<Self>,
     ) -> bool {
-        if self
+        let tracked_child_pane = self.child_agent_panes.get(&child_conversation_id).copied();
-            .child_agent_panes
+        match tracked_child_pane {
-            .get(&child_conversation_id)
+            Some(pane_id) if self.has_pane_id(pane_id) => return true,
-            .is_some_and(|pane_id| self.has_pane_id(*pane_id))
+            Some(pane_id) => {
-        {
+                log::warn!(
-            return true;
+                    "Repairing child conversation {child_conversation_id:?} with missing pane \
+                     {pane_id:?}"
+                );
+                self.child_agent_panes.remove(&child_conversation_id);
+            }
+            None => {}
         }
+        let had_stale_mapping = tracked_child_pane.is_some();
 
         let parent_conversation_id =
             BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
@@ -136,18 +142,42 @@ impl PaneGroup {
             });
 
         let Some(parent_conversation_id) = parent_conversation_id else {
-            return self
+            return !had_stale_mapping
-                .terminal_view_id_for_owned_conversation(child_conversation_id, ctx)
+                && self
-                .is_some();
+                    .terminal_view_id_for_owned_conversation(child_conversation_id, ctx)
+                    .is_some();
         };
 
         let child_owner_terminal_view_id =
             self.terminal_view_id_for_owned_conversation(child_conversation_id, ctx);
         let Some(parent_pane_id) = self.pane_id_for_owned_conversation(parent_conversation_id, ctx)
         else {
-            return child_owner_terminal_view_id.is_some();
+            return !had_stale_mapping && child_owner_terminal_view_id.is_some();
         };
 
+        if had_stale_mapping {
+            let child_conversation = BlocklistAIHistoryModel::as_ref(ctx)
+                .conversation(&child_conversation_id)
+                .cloned()
+                .or_else(|| {
+                    RestoredAgentConversations::handle(ctx).update(ctx, |store, _| {
+                        store.take_conversation(&child_conversation_id)
+                    })
+                });
+            let Some(child_conversation) = child_conversation else {
+                log::warn!(
+                    "Cannot repair missing pane for child conversation \
+                     {child_conversation_id:?}: conversation not found"
+                );
+                return false;
+            };
+            self.create_hidden_child_agent_pane(child_conversation, parent_pane_id, ctx);
+            return self
+                .child_agent_panes
+                .get(&child_conversation_id)
+                .is_some_and(|pane_id| self.has_pane_id(*pane_id));
+        }
+
         if self.is_conversation_owned_outside_pane(child_conversation_id, parent_pane_id, ctx) {
             return true;
         }

app/src/pane_group/mod.rs

@@ -4692,12 +4692,29 @@ impl PaneGroup {
         // transfer-on-close step below.
         if self.is_child_agent_pane(pane_id) {
             // Revert the swap if the child is currently swapped in.
-            if self.panes.original_pane_for_replacement(pane_id).is_some() {
+            if let Some(original_pane_id) = self.panes.original_pane_for_replacement(pane_id) {
-                self.panes.revert_temporary_replacement(pane_id);
+                if self.panes.revert_temporary_replacement(pane_id) != Some(original_pane_id) {
+                    log::error!("close_pane: failed to restore pane replaced by child {pane_id:?}");
+                    return;
+                }
+                self.remove_from_pane_history(pane_id);
+                self.focus_pane_preserving_maximized_state(original_pane_id, true, ctx);
+                self.update_pane_history(original_pane_id);
             }
             // Or remove the child from the tree if it was split off.
-            else if self.panes.is_pane_in_tree(pane_id) && !self.panes.remove(pane_id) {
+            else if self.panes.is_pane_in_tree(pane_id) {
-                log::error!("close_pane: failed to remove split-off child pane from tree");
+                // Focus must move while the child is still in the tree so
+                // pane ordering can select a valid sibling.
+                self.focus_next_terminal_pane_and_activate_session(
+                    pane_id,
+                    PaneRemovalReason::Close,
+                    ctx,
+                );
+                if !self.panes.remove(pane_id) {
+                    log::error!("close_pane: failed to remove split-off child pane from tree");
+                }
+            } else {
+                self.remove_from_pane_history(pane_id);
             }
             // Drop any leftover swap entry recording this child as the
             // original side. Otherwise a later revert of the surviving
@@ -4711,13 +4728,6 @@ impl PaneGroup {
                     view.clear_orchestration_split_off(ctx);
                 });
             }
-            self.panes.remove_hidden_pane(pane_id);
-            self.focus_next_terminal_pane_and_activate_session(
-                pane_id,
-                PaneRemovalReason::Close,
-                ctx,
-            );
-            self.pane_contents.remove(&pane_id);
             self.handle_pane_count_change(ctx);
             ctx.emit(Event::TerminalViewStateChanged);
             ctx.emit(Event::AppStateChanged);
@@ -6948,7 +6958,11 @@ impl PaneGroup {
         conversation_id: AIConversationId,
         ctx: &mut ViewContext<Self>,
     ) {
-        let from_child_panes = self.child_agent_panes.get(&conversation_id).copied();
+        let from_child_panes = self
+            .child_agent_panes
+            .get(&conversation_id)
+            .copied()
+            .filter(|pane_id| self.has_pane_id(*pane_id));
         let from_visible_pane = self
             .find_visible_terminal_pane_for_conversation(conversation_id, ctx)
             .map(PaneId::from);
@@ -7069,6 +7083,14 @@ impl PaneGroup {
         ctx: &mut ViewContext<Self>,
     ) -> Option<PaneId> {
         let child_pane_id = self.child_agent_panes.get(&conversation_id).copied()?;
+        if !self.has_pane_id(child_pane_id) {
+            log::error!(
+                "unhide_child_agent_pane_for_split_off: child conversation {conversation_id:?} \
+                 references missing pane {child_pane_id:?}"
+            );
+            self.child_agent_panes.remove(&conversation_id);
+            return None;
+        }
 
         // If the child was previously split off and then swapped over,
         // it's recorded as the original of an active swap. Revert that

app/src/pane_group/mod_tests.rs

@@ -783,6 +783,66 @@ fn test_insert_hidden_child_agent_pane_keeps_focus_and_active_session() {
     });
 }
 
+#[test]
+fn test_closing_split_off_child_keeps_backing_pane_for_reopen() {
+    App::test((), |mut app| async move {
+        initialize_app(&mut app);
+        let pane_group = mock_pane_group(&mut app, Default::default());
+
+        pane_group.update(&mut app, |panes, ctx| {
+            let parent_pane_id = get_newly_created_pane_id(panes, &[]);
+            let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx);
+            let child = create_hidden_child_agent_conversation(
+                panes,
+                HiddenChildAgentConversationRequest {
+                    parent_pane_id,
+                    name: "Agent 1".to_string(),
+                    parent_conversation_id,
+                    orchestration_harness: None,
+                    env_vars: HashMap::new(),
+                    task_context: None,
+                    is_shared_session_creator: IsSharedSessionCreator::No,
+                },
+                ctx,
+            )
+            .expect("fresh hidden child conversation should be created");
+            let child_pane_id = panes
+                .child_agent_panes
+                .get(&child.conversation_id)
+                .copied()
+                .expect("fresh hidden child pane should be tracked");
+
+            assert_eq!(
+                panes.unhide_child_agent_pane_for_split_off(child.conversation_id, ctx),
+                Some(child_pane_id)
+            );
+            assert!(panes.panes.is_pane_in_tree(child_pane_id));
+
+            panes.close_pane(child_pane_id, ctx);
+
+            assert_eq!(panes.focused_pane_id(ctx), parent_pane_id);
+            assert!(panes.has_pane_id(child_pane_id));
+            assert_eq!(
+                panes.child_agent_panes.get(&child.conversation_id),
+                Some(&child_pane_id)
+            );
+            assert!(!panes.panes.is_pane_in_tree(child_pane_id));
+
+            assert_eq!(
+                panes.unhide_child_agent_pane_for_split_off(child.conversation_id, ctx),
+                Some(child_pane_id)
+            );
+            assert_eq!(panes.focused_pane_id(ctx), child_pane_id);
+            assert!(panes.panes.is_pane_in_tree(child_pane_id));
+            assert!(panes
+                .panes
+                .pane_ids()
+                .into_iter()
+                .all(|pane_id| panes.has_pane_id(pane_id)));
+        });
+    });
+}
+
 #[test]
 fn test_swapping_to_child_agent_from_maximized_pane_keeps_maximized_state() {
     App::test((), |mut app| async move {

crates/galaxy_agent_rig/src/stream.rs

@@ -395,9 +395,28 @@ fn completion_error_indicates_recoverable_transport(error: &CompletionError) ->
     }
 }
 
+fn completion_error_indicates_transient_provider_failure(error: &CompletionError) -> bool {
+    error
+        .provider_response_body()
+        .is_some_and(text_indicates_transient_provider_failure)
+        || text_indicates_transient_provider_failure(&error.to_string())
+}
+
+fn text_indicates_transient_provider_failure(text: &str) -> bool {
+    let normalized = text.to_ascii_lowercase();
+    normalized.contains("server_is_overloaded")
+        || normalized.contains("service_unavailable_error")
+        || normalized.contains("temporarily unavailable")
+        || normalized.contains("service unavailable")
+        || normalized.contains("server is overloaded")
+        || normalized.contains("servers are currently overloaded")
+}
+
 fn map_completion_error(error: CompletionError) -> AgentError {
     let is_context_window_exceeded = completion_error_indicates_context_window_exceeded(&error);
     let is_recoverable_transport = completion_error_indicates_recoverable_transport(&error);
+    let is_transient_provider_failure =
+        completion_error_indicates_transient_provider_failure(&error);
     let status = error
         .provider_response_status()
         .map(|status| status.as_u16());
@@ -430,7 +449,8 @@ fn map_completion_error(error: CompletionError) -> AgentError {
     mapped.recoverable = matches!(
         kind,
         AgentErrorKind::RateLimited | AgentErrorKind::Transport
-    ) || status.is_some_and(|status| (500..=599).contains(&status));
+    ) || status.is_some_and(|status| (500..=599).contains(&status))
+        || is_transient_provider_failure;
     mapped
 }
 

crates/galaxy_agent_rig/src/stream_tests.rs

@@ -14,3 +14,16 @@ fn flattened_sse_http_client_error_is_recoverable_transport() {
     assert_eq!(mapped.kind, AgentErrorKind::Transport);
     assert!(mapped.recoverable);
 }
+
+#[test]
+fn flattened_streamed_provider_overload_is_recoverable() {
+    let error = CompletionError::ProviderError(
+        r#"{"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}"#
+            .to_string(),
+    );
+
+    let mapped = map_completion_error(error);
+
+    assert_eq!(mapped.kind, AgentErrorKind::Provider);
+    assert!(mapped.recoverable);
+}