magi-code 0.96.1

use super::*;

fn join_title_if_finished(title_handle: &mut Option<std::thread::JoinHandle<()>>) -> bool {
    if title_handle
        .as_ref()
        .is_some_and(std::thread::JoinHandle::is_finished)
    {
        let _ = title_handle.take().expect("checked title handle").join();
        true
    } else {
        false
    }
}

pub(super) fn finish_title_generation_for_mode(
    title_handle: &mut Option<std::thread::JoinHandle<()>>,
    invocation_mode: InvocationMode,
) {
    match invocation_mode {
        InvocationMode::Print => {
            if !join_title_if_finished(title_handle) {
                let _ = title_handle.take();
            }
        }
        InvocationMode::Side | InvocationMode::LocalAgent => {
            if let Some(handle) = title_handle.take() {
                let _ = handle.join();
            }
        }
        InvocationMode::MissionControl | InvocationMode::Subagent => {
            join_title_if_finished(title_handle);
        }
    }
}

pub(super) struct TitleGenerationGuard {
    handle: Option<std::thread::JoinHandle<()>>,
    pub(super) invocation_mode: InvocationMode,
    pub(super) cancellation: AgentCancellation,
}

impl TitleGenerationGuard {
    pub(super) fn new(
        handle: Option<std::thread::JoinHandle<()>>,
        invocation_mode: InvocationMode,
        cancellation: AgentCancellation,
    ) -> Self {
        Self {
            handle,
            invocation_mode,
            cancellation,
        }
    }

    pub(super) fn finish(&mut self) {
        finish_title_generation_for_mode(&mut self.handle, self.invocation_mode);
    }
}

impl Drop for TitleGenerationGuard {
    fn drop(&mut self) {
        if !self.cancellation.is_canceled() || self.invocation_mode == InvocationMode::LocalAgent {
            self.finish();
        }
    }
}