IdeA/crates/application/src/workstate/actions.rs

//! Agent-level controlled actions over live sessions (Lot D).
//!
//! Companion to the read-only work-state model: these use cases let the UI drive
//! an **already-running** agent's live session by **agent id**, never by spawning a
//! new process. [`AttachLiveAgent`] rebinds the hosting cell ("a cell is a view");
//! [`StopLiveAgent`] tears a PTY live session down or cancels the current structured
//! turn while preserving the structured session. Neither touches the agent entity, its tickets, the
//! conversation summaries, the handoff or the provider sessions.

use std::sync::Arc;

use async_trait::async_trait;
use domain::input::InputMediator;
use domain::{AgentId, NodeId, Project, RuntimeAgentKey, SessionId};

use crate::error::AppError;
use crate::orchestrator::OrchestratorService;
use crate::terminal::{
    stop_live_agent_runtime_in_project, CloseTerminal, LiveRuntimeState, LiveSessionKind,
    LiveSessions,
};

/// Input for [`AttachLiveAgent::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttachLiveAgentInput {
    /// The owning project (boundary/symmetry; the registries are global by agent).
    pub project: Project,
    /// The already-running agent whose live session is rebound.
    pub agent_id: AgentId,
    /// The layout leaf that should host the live session view.
    pub node_id: NodeId,
}

/// Output of [`AttachLiveAgent::execute`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AttachLiveAgentOutput {
    /// The rebound agent.
    pub agent_id: AgentId,
    /// The node now hosting the session view.
    pub node_id: NodeId,
    /// The unchanged live session id.
    pub session_id: SessionId,
    /// Runtime family that owns the session.
    pub kind: LiveSessionKind,
}

/// Rebinds an already-running agent session to a visible layout cell, never
/// spawning a process ("a cell is a view").
///
/// Resolves the live session by agent in the PTY registry first, then the
/// structured one, and updates only the hosting `node_id` (the PTY handle / agent
/// session, its id and scrollback stay intact). Idempotent when the agent is
/// already attached to the same node (the rebind is a no-op that returns the same
/// session). No log/handoff/provider-session is touched.
pub struct AttachLiveAgent {
    live: Arc<LiveSessions>,
}

impl AttachLiveAgent {
    /// Builds the use case from the aggregated live-session registry.
    #[must_use]
    pub fn new(live: Arc<LiveSessions>) -> Self {
        Self { live }
    }

    /// Rebinds the agent's live session onto `node_id`.
    ///
    /// # Errors
    /// [`AppError::NotFound`] when the agent has no live session in either registry.
    pub fn execute(&self, input: AttachLiveAgentInput) -> Result<AttachLiveAgentOutput, AppError> {
        // PTY first, then structured (one-live-session-per-agent ⇒ at most one match).
        if let Some(session) = self.live.pty.rebind_agent_node_in_project(
            input.project.id,
            &input.agent_id,
            input.node_id,
        ) {
            return Ok(AttachLiveAgentOutput {
                agent_id: input.agent_id,
                node_id: session.node_id,
                session_id: session.id,
                kind: LiveSessionKind::Pty,
            });
        }
        if let Some(session) = self.live.structured.rebind_agent_node_in_project(
            input.project.id,
            &input.agent_id,
            input.node_id,
        ) {
            return Ok(AttachLiveAgentOutput {
                agent_id: input.agent_id,
                node_id: input.node_id,
                session_id: session.id(),
                kind: LiveSessionKind::Structured,
            });
        }
        Err(AppError::NotFound(format!(
            "running session for agent {}",
            input.agent_id
        )))
    }
}

/// Input for [`StopLiveAgent::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StopLiveAgentInput {
    /// The owning project (boundary/symmetry; the registries are global by agent).
    pub project: Project,
    /// The agent whose live session is stopped.
    pub agent_id: AgentId,
}

/// Output of [`StopLiveAgent::execute`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StopLiveAgentOutput {
    /// The stopped agent.
    pub agent_id: AgentId,
    /// The id of the session that was torn down.
    pub session_id: SessionId,
    /// Last known layout node hosting the runtime view.
    pub node_id: NodeId,
    /// Runtime family that owned the session.
    pub kind: LiveSessionKind,
    /// Runtime lifecycle projection after the stop request.
    pub runtime_state: LiveRuntimeState,
}

/// Best-effort stopper for execution resources owned by one live agent turn.
///
/// Implementations must be idempotent: a missing, already-finished or already-stopped
/// child resource is a successful no-op. Failures are deliberately not propagated by
/// [`StopLiveAgent`], because the product contract is "stop the agent; cascade related
/// work best-effort".
#[async_trait]
pub trait StopAgentExecutionCascade: Send + Sync {
    /// Requests cancellation of child execution resources for `agent_id` in `project`.
    async fn stop_agent_execution_children(&self, project: &Project, agent_id: AgentId);
}

/// Stops an already-running agent's current work.
///
/// Resolves the session by agent (structured first, then PTY). A structured
/// session keeps its registry entry and receives `cancel_current_turn()`, so the
/// provider conversation/context survives the stop. A PTY session has no finer
/// turn boundary, so it is torn down through the existing [`CloseTerminal`]
/// primitive (registry removal + process kill).
pub struct StopLiveAgent {
    live: Arc<LiveSessions>,
    close: Arc<CloseTerminal>,
    input: Option<Arc<dyn InputMediator>>,
    waits: Option<Arc<OrchestratorService>>,
    cascade: Option<Arc<dyn StopAgentExecutionCascade>>,
}

impl StopLiveAgent {
    /// Builds the use case from the live-session registry and the close primitive.
    #[must_use]
    pub fn new(live: Arc<LiveSessions>, close: Arc<CloseTerminal>) -> Self {
        Self {
            live,
            close,
            input: None,
            waits: None,
            cascade: None,
        }
    }

    /// Wires cancellation helpers used for user-driven stop cascades.
    ///
    /// Stopping A while A waits on B should stop B too. The orchestrator service is
    /// only used as a read-only provider of active wait edges; the input mediator is
    /// used to preempt the running turn before the session is torn down.
    #[must_use]
    pub fn with_cascade(
        mut self,
        input: Arc<dyn InputMediator>,
        waits: Arc<OrchestratorService>,
    ) -> Self {
        self.input = Some(input);
        self.waits = Some(waits);
        self
    }

    /// Wires best-effort cancellation of child resources spawned by the running turn.
    #[must_use]
    pub fn with_execution_cascade(mut self, cascade: Arc<dyn StopAgentExecutionCascade>) -> Self {
        self.cascade = Some(cascade);
        self
    }

    /// Stops the agent's current work.
    ///
    /// # Errors
    /// - [`AppError::NotFound`] when the agent has no live session,
    /// - [`AppError::Process`] when the PTY kill or structured shutdown fails.
    pub async fn execute(
        &self,
        input: StopLiveAgentInput,
    ) -> Result<StopLiveAgentOutput, AppError> {
        self.stop_dependencies(&input.project, input.agent_id)
            .await?;
        if let Some(mediator) = &self.input {
            mediator.preempt(RuntimeAgentKey::new(input.project.id, input.agent_id));
        }
        if let Some(out) = self
            .cancel_structured_turn_in_project(&input.project, input.agent_id)
            .await?
        {
            return Ok(out);
        }
        self.stop_one_in_project(&input.project, input.agent_id)
            .await
    }

    async fn stop_dependencies(
        &self,
        project: &Project,
        agent_id: AgentId,
    ) -> Result<(), AppError> {
        if let Some(cascade) = &self.cascade {
            cascade
                .stop_agent_execution_children(project, agent_id)
                .await;
        }
        let Some(waits) = &self.waits else {
            return Ok(());
        };
        for dep in waits.active_wait_dependencies(agent_id) {
            if let Some(cascade) = &self.cascade {
                cascade.stop_agent_execution_children(project, dep).await;
            }
            if let Some(mediator) = &self.input {
                mediator.preempt(RuntimeAgentKey::new(project.id, dep));
            }
            if self
                .cancel_structured_turn_in_project(project, dep)
                .await?
                .is_none()
            {
                let _ = self.stop_one_in_project(project, dep).await;
            }
        }
        Ok(())
    }

    async fn cancel_structured_turn_in_project(
        &self,
        project: &Project,
        agent_id: AgentId,
    ) -> Result<Option<StopLiveAgentOutput>, AppError> {
        let Some(snapshot) = self
            .live
            .live_agent_snapshots_for(project.id, &agent_id)
            .into_iter()
            .find(|snapshot| snapshot.kind == LiveSessionKind::Structured)
        else {
            return Ok(None);
        };
        let Some(session) = self
            .live
            .structured
            .session_for_agent_in_project(project.id, &agent_id)
        else {
            return Ok(None);
        };
        self.live
            .mark_runtime_state(snapshot, LiveRuntimeState::Stopping);
        session
            .cancel_current_turn()
            .await
            .map_err(|err| AppError::Process(err.to_string()))?;
        self.live
            .mark_runtime_state(snapshot, LiveRuntimeState::Stopped);
        Ok(Some(StopLiveAgentOutput {
            agent_id: snapshot.agent_id,
            session_id: snapshot.session_id,
            node_id: snapshot.node_id,
            kind: snapshot.kind,
            runtime_state: LiveRuntimeState::Stopped,
        }))
    }

    async fn stop_one_in_project(
        &self,
        project: &Project,
        agent_id: AgentId,
    ) -> Result<StopLiveAgentOutput, AppError> {
        stop_live_agent_runtime_in_project(project, agent_id, &self.live, &self.close)
            .await?
            .into_iter()
            .next()
            .map(|stopped| StopLiveAgentOutput {
                agent_id: stopped.agent_id,
                session_id: stopped.session_id,
                node_id: stopped.node_id,
                kind: stopped.kind,
                runtime_state: stopped.state,
            })
            .or_else(|| {
                self.live
                    .runtime_status_for(project.id, &agent_id)
                    .filter(|status| {
                        matches!(
                            status.state,
                            LiveRuntimeState::Stopping | LiveRuntimeState::Stopped
                        )
                    })
                    .map(|status| StopLiveAgentOutput {
                        agent_id: status.agent_id,
                        session_id: status.session_id,
                        node_id: status.node_id,
                        kind: status.kind,
                        runtime_state: status.state,
                    })
            })
            .ok_or_else(|| AppError::NotFound(format!("running session for agent {agent_id}")))
    }
}