magi-code 0.80.1

use crate::{
    agent::steering::AgentSteering,
    cancellation::{AgentCancellation, AgentCancellationHandle},
    output::ActivityId,
};
use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
};

const MAX_RETAINED_SUBAGENT_CONTROLS: usize = 64;

/// One live parent run's children, keyed by raw runtime activity IDs.
#[derive(Clone, Debug, Default)]
pub(crate) struct SubagentControls {
    inner: Arc<Mutex<HashMap<ActivityId, SubagentControl>>>,
}

#[derive(Clone, Debug)]
pub(crate) struct SubagentControl {
    pub(crate) steering: AgentSteering,
    cancellation: AgentCancellationHandle,
    cancellation_token: AgentCancellation,
    active: Arc<AtomicBool>,
}

impl SubagentControl {
    pub(crate) fn cancel(&self) {
        self.cancellation.cancel();
        self.steering.close();
    }

    pub(crate) fn is_active(&self) -> bool {
        self.active.load(Ordering::SeqCst) && !self.cancellation_token.is_canceled()
    }
}

impl SubagentControls {
    pub(crate) fn has_pending_input(&self) -> bool {
        self.inner
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .values()
            .any(|control| control.steering.try_pending_count() != Some(0))
    }

    pub(crate) fn get(&self, id: &ActivityId) -> Option<SubagentControl> {
        self.inner
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .get(id)
            .cloned()
    }

    pub(super) fn register(
        &self,
        id: ActivityId,
        steering: AgentSteering,
        cancellation: AgentCancellationHandle,
        cancellation_token: AgentCancellation,
    ) -> Option<SubagentControlRegistration> {
        let mut entries = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        entries.retain(|_, control| {
            control.active.load(Ordering::SeqCst) || control.steering.try_pending_count() != Some(0)
        });
        // Refuse new controls rather than evict accepted, unacknowledged input.
        if entries.len() >= MAX_RETAINED_SUBAGENT_CONTROLS || entries.contains_key(&id) {
            return None;
        }
        let control = SubagentControl {
            steering,
            cancellation,
            cancellation_token,
            active: Arc::new(AtomicBool::new(true)),
        };
        entries.insert(id.clone(), control.clone());
        Some(SubagentControlRegistration {
            registry: self.clone(),
            id,
            control,
        })
    }
}

/// Closes stale UI handles, retaining unacknowledged input after any exit.
pub(super) struct SubagentControlRegistration {
    registry: SubagentControls,
    id: ActivityId,
    control: SubagentControl,
}

impl Drop for SubagentControlRegistration {
    fn drop(&mut self) {
        self.control.steering.close();
        self.control.active.store(false, Ordering::SeqCst);
        let has_pending_input = self.control.steering.pending_count() > 0;
        let mut entries = self
            .registry
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if entries
            .get(&self.id)
            .is_some_and(|entry| Arc::ptr_eq(&entry.active, &self.control.active))
            && !has_pending_input
        {
            entries.remove(&self.id);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn retained_input_is_not_evicted_at_capacity_and_recall_frees_admission() {
        let controls = SubagentControls::default();
        for index in 0..MAX_RETAINED_SUBAGENT_CONTROLS {
            let (token, handle) = AgentCancellation::default().child_token();
            let steering = AgentSteering::new();
            steering.try_enqueue(format!("pending {index}")).unwrap();
            let registration = controls
                .register(ActivityId::new(index.to_string()), steering, handle, token)
                .unwrap();
            drop(registration);
        }
        let (token, handle) = AgentCancellation::default().child_token();
        assert!(
            controls
                .register(
                    ActivityId::new("overflow"),
                    AgentSteering::new(),
                    handle.clone(),
                    token.clone()
                )
                .is_none()
        );
        let retained = controls.get(&ActivityId::new("0")).unwrap();
        assert!(!retained.is_active());
        assert_eq!(
            retained.steering.restore_pending_messages(|text| {
                assert_eq!(text, "pending 0");
                true
            }),
            Ok(true)
        );
        assert!(
            controls
                .register(ActivityId::new("next"), AgentSteering::new(), handle, token)
                .is_some()
        );
        assert!(controls.get(&ActivityId::new("1")).is_some());
    }
}

#[cfg(test)]
#[test]
fn cancel_signals_before_busy_persistence_finishes() {
    let controls = SubagentControls::default();
    let (token, handle) = AgentCancellation::default().child_token();
    let steering = AgentSteering::new();
    steering.try_enqueue("recover me".into()).unwrap();
    let registration = controls
        .register(
            ActivityId::new("busy"),
            steering.clone(),
            handle,
            token.clone(),
        )
        .unwrap();
    let control = controls.get(&ActivityId::new("busy")).unwrap();
    let batch = steering.observe_collapsed().unwrap();
    let result = steering.persist_and_acknowledge(&batch, || {
        control.cancel();
        assert!(token.is_canceled());
        assert!(!control.is_active());
        Err::<(), _>("failed write")
    });
    assert_eq!(result, Err("failed write"));
    drop(registration);
    assert!(controls.has_pending_input());
    assert_eq!(
        controls
            .get(&ActivityId::new("busy"))
            .unwrap()
            .steering
            .restore_pending_messages(|text| {
                assert_eq!(text, "recover me");
                true
            }),
        Ok(true)
    );
}