IdeA/crates/application/src/orchestrator/wake.rs

//! Owner wake service for first-class background task completions.
//!
//! This application adapter implements the domain [`AgentWakePort`] above the
//! existing structured/headless session path. It owns no concrete process or
//! Tauri wiring: B7 injects the session provider that reuses
//! `StructuredSessions`/`LaunchAgent`.

use std::sync::Arc;

use async_trait::async_trait;
use domain::background_task::{BackgroundTask, BackgroundTaskResult};
use domain::events::DomainEvent;
use domain::ids::{AgentId, RuntimeAgentKey};
use domain::inbox::{AgentInbox, InboxItem, InboxItemKind, InboxSource};
use domain::input::InputMediator;
use domain::mailbox::{AgentMailbox, Ticket};
use domain::ports::{
    AgentSession, AgentWakePort, BackgroundTaskStore, EventBus, WakeError, WakeReason,
};
use domain::project::Project;

use crate::agent::drain_reply_stream_with_readiness;

/// Provides a structured/headless session for a wake turn.
///
/// The production implementation is expected to return the live
/// `StructuredSessions` entry when present, or launch/reattach the agent using the
/// same mechanism as inter-agent delegation when absent.
#[async_trait]
pub trait WakeSessionProvider: Send + Sync {
    /// Returns or starts the structured session used to send the wake prompt.
    ///
    /// # Errors
    /// [`WakeError::Session`] when the agent cannot be driven headlessly.
    async fn session_for_wake(
        &self,
        project: &Project,
        agent: AgentId,
    ) -> Result<Arc<dyn AgentSession>, WakeError>;
}

/// Drain-then-wake implementation of [`AgentWakePort`].
pub struct AgentWakeService {
    inbox: Arc<dyn AgentInbox>,
    input: Arc<dyn InputMediator>,
    mailbox: Arc<dyn AgentMailbox>,
    tasks: Arc<dyn BackgroundTaskStore>,
    sessions: Arc<dyn WakeSessionProvider>,
    events: Option<Arc<dyn EventBus>>,
}

impl AgentWakeService {
    /// Builds a wake service from its ports.
    #[must_use]
    pub fn new(
        inbox: Arc<dyn AgentInbox>,
        input: Arc<dyn InputMediator>,
        mailbox: Arc<dyn AgentMailbox>,
        tasks: Arc<dyn BackgroundTaskStore>,
        sessions: Arc<dyn WakeSessionProvider>,
        events: Option<Arc<dyn EventBus>>,
    ) -> Self {
        Self {
            inbox,
            input,
            mailbox,
            tasks,
            sessions,
            events,
        }
    }

    fn publish(&self, event: DomainEvent) {
        if let Some(events) = &self.events {
            events.publish(event);
        }
    }

    fn publish_failed(&self, project: &Project, agent: AgentId, reason: impl ToString) {
        self.publish(DomainEvent::AgentWakeFailed {
            project_id: project.id,
            agent_id: agent,
            reason: reason.to_string(),
        });
    }

    async fn wake_one(
        &self,
        project: &Project,
        agent: AgentId,
        reason: WakeReason,
    ) -> Result<(), WakeError> {
        let key = RuntimeAgentKey::new(project.id, agent);
        self.publish(DomainEvent::AgentWakeScheduled {
            project_id: project.id,
            agent_id: agent,
        });

        if self.input.busy_state(key).is_busy() {
            return Err(WakeError::AgentBusy { agent_id: agent });
        }

        let Some(item) = self.inbox.dequeue_next(key) else {
            return Ok(());
        };
        let delivery = self.delivery_from_item(item, &reason).await?;
        let ticket = Ticket::new(delivery.ticket_id, "IdeA", delivery.prompt.clone());
        let _pending = self.input.enqueue_silent(key, ticket);
        let guard = WakeTurnGuard::new(
            Arc::clone(&self.input),
            Arc::clone(&self.mailbox),
            key,
            delivery.ticket_id,
        );

        self.publish(DomainEvent::AgentWakeStarted {
            project_id: project.id,
            agent_id: agent,
        });
        let session = self.sessions.session_for_wake(project, agent).await?;
        let stream = session
            .send(&delivery.prompt)
            .await
            .map_err(|err| WakeError::Session(err.to_string()))?;

        if let Some(task_id) = delivery.delivered_task_id {
            self.tasks
                .mark_completion_delivered(task_id)
                .await
                .map_err(|err| WakeError::Store(err.to_string()))?;
            self.publish(DomainEvent::BackgroundTaskCompletionDelivered {
                project_id: project.id,
                task_id,
                owner_agent_id: agent,
            });
        }
        let final_content = drain_reply_stream_with_readiness(stream, self.input.as_ref(), key)
            .await
            .map_err(|err| WakeError::Session(err.to_string()))?;

        if let Some(task_id) = delivery.delivered_task_id {
            let task = self.load_task(task_id).await?;
            if let Some(link) = task.rendezvous.as_ref() {
                let still_waiting = self
                    .tasks
                    .list_open_for_agent(agent)
                    .await
                    .map_err(|err| WakeError::Store(err.to_string()))?
                    .into_iter()
                    .any(|open| {
                        open.id != task_id
                            && open.rendezvous.as_ref().is_some_and(|open_link| {
                                open_link.rendezvous_id == link.rendezvous_id
                            })
                    });
                if !still_waiting {
                    self.mailbox
                        .resolve_ticket(key, link.ticket_id, final_content)
                        .map_err(|err| WakeError::Task(err.to_string()))?;
                }
            }
        }

        self.mailbox.cancel_head(key, delivery.ticket_id);
        self.input.mark_idle(key);
        guard.disarm();
        Ok(())
    }

    async fn delivery_from_item(
        &self,
        item: InboxItem,
        reason: &WakeReason,
    ) -> Result<WakeDelivery, WakeError> {
        match (reason, item.kind, item.source) {
            (
                WakeReason::BackgroundCompletion { task_id },
                InboxItemKind::BackgroundCompletion,
                InboxSource::BackgroundTask {
                    task_id: item_task_id,
                },
            ) if *task_id == item_task_id => {
                let task = self.load_task(item_task_id).await?;
                Ok(WakeDelivery {
                    ticket_id: item.id,
                    prompt: background_completion_prompt(&task),
                    delivered_task_id: Some(item_task_id),
                })
            }
            (WakeReason::InboxItem { item_id }, _, _) if *item_id == item.id => Ok(WakeDelivery {
                ticket_id: item.id,
                prompt: system_prompt_from_item(&item),
                delivered_task_id: None,
            }),
            _ => Err(WakeError::UnexpectedInboxItem(format!(
                "item={} kind={:?} source={:?} reason={:?}",
                item.id, item.kind, item.source, reason
            ))),
        }
    }

    async fn load_task(&self, task_id: domain::ids::TaskId) -> Result<BackgroundTask, WakeError> {
        self.tasks
            .get(task_id)
            .await
            .map_err(|err| WakeError::Store(err.to_string()))?
            .ok_or_else(|| WakeError::Task(format!("background task {task_id} not found")))
            .and_then(|task| {
                if task.result.is_some() {
                    Ok(task)
                } else {
                    Err(WakeError::Task(format!(
                        "background task {task_id} has no terminal result"
                    )))
                }
            })
    }
}

#[async_trait]
impl AgentWakePort for AgentWakeService {
    async fn wake_agent(
        &self,
        project: &Project,
        agent: AgentId,
        reason: WakeReason,
    ) -> Result<(), WakeError> {
        match self.wake_one(project, agent, reason).await {
            Ok(()) => Ok(()),
            Err(err) => {
                self.publish_failed(project, agent, &err);
                Err(err)
            }
        }
    }
}

struct WakeDelivery {
    ticket_id: domain::mailbox::TicketId,
    prompt: String,
    delivered_task_id: Option<domain::ids::TaskId>,
}

struct WakeTurnGuard {
    input: Arc<dyn InputMediator>,
    mailbox: Arc<dyn AgentMailbox>,
    agent: RuntimeAgentKey,
    ticket: domain::mailbox::TicketId,
    armed: bool,
}

impl WakeTurnGuard {
    fn new(
        input: Arc<dyn InputMediator>,
        mailbox: Arc<dyn AgentMailbox>,
        agent: RuntimeAgentKey,
        ticket: domain::mailbox::TicketId,
    ) -> Self {
        Self {
            input,
            mailbox,
            agent,
            ticket,
            armed: true,
        }
    }

    fn disarm(mut self) {
        self.armed = false;
    }
}

impl Drop for WakeTurnGuard {
    fn drop(&mut self) {
        if self.armed {
            self.mailbox.cancel_head(self.agent, self.ticket);
            self.input.mark_idle(self.agent);
        }
    }
}

fn system_prompt_from_item(item: &InboxItem) -> String {
    format!("Message système pour l'agent :\n\n{}", item.body)
}

fn background_completion_prompt(task: &BackgroundTask) -> String {
    match task.result.as_ref() {
        Some(BackgroundTaskResult::Success {
            exit_code,
            summary,
            stdout_tail,
            stderr_tail,
            ..
        }) => format!(
            "La tâche de fond {} est terminée : succès{}.\nRésumé : {}\n{}{}",
            task.id,
            exit_code_text(*exit_code),
            summary,
            tail_text("stdout", stdout_tail.as_deref()),
            tail_text("stderr", stderr_tail.as_deref())
        ),
        Some(BackgroundTaskResult::Failure {
            exit_code,
            error,
            stdout_tail,
            stderr_tail,
            ..
        }) => format!(
            "La tâche de fond {} est terminée : échec{}.\nErreur : {}\n{}{}",
            task.id,
            exit_code_text(*exit_code),
            error,
            tail_text("stdout", stdout_tail.as_deref()),
            tail_text("stderr", stderr_tail.as_deref())
        ),
        Some(BackgroundTaskResult::Cancelled { reason, .. }) => format!(
            "La tâche de fond {} est terminée : annulée.\nRaison : {}",
            task.id, reason
        ),
        Some(BackgroundTaskResult::CancelFailed { error, .. }) => format!(
            "La tâche de fond {} est terminée : échec de l'annulation.\nErreur : {}",
            task.id, error
        ),
        Some(BackgroundTaskResult::Expired { reason, .. }) => format!(
            "La tâche de fond {} est terminée : expirée.\nRaison : {}",
            task.id, reason
        ),
        None => format!(
            "La tâche de fond {} est terminée, mais son résultat est indisponible.",
            task.id
        ),
    }
}

fn exit_code_text(exit_code: Option<i32>) -> String {
    exit_code.map_or_else(String::new, |code| format!(" (exit {code})"))
}

fn tail_text(label: &str, value: Option<&str>) -> String {
    value.map_or_else(String::new, |text| {
        if text.is_empty() {
            String::new()
        } else {
            format!("{label} :\n{text}\n")
        }
    })
}