maestro/crates/maestro-tui/src/structured_task.rs

//! Live projection state for one structured session.

use std::collections::HashSet;
use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use futures_util::StreamExt;
use maestro_protocol::{
    ReducedState, StructuredApprovalDecision, StructuredEvent, TranscriptDelta, TranscriptRow,
    TranscriptRowKind,
};
use tokio::sync::watch;
use tokio_tungstenite::tungstenite::{Error as WsError, Message};

use crate::input::InputKey;
use crate::session_task::SessionState;
use crate::text_input::TextInput;

/// Renderer and input state built directly from the daemon's two projections.
#[derive(Debug, Clone)]
pub struct StructuredState {
    pub rows: Vec<TranscriptRow>,
    pub reduced: ReducedState,
    pub selected: Option<String>,
    pub expanded: HashSet<String>,
    pub scroll: usize,
    /// Row id per rendered line from the last render pass, parallel to that
    /// pass's `lines`. Lets the next render translate `scroll` — a raw index
    /// into that now-stale layout — back into a row identity before mapping
    /// it onto the freshly rebuilt layout, so growth/shrinkage above the
    /// viewport shifts indices without shifting what's on screen.
    pub(crate) line_rows: Vec<Option<String>>,
    pub follow_selection: bool,
    pub composer: TextInput,
    pub error: Option<String>,
    submitting: bool,
    submission_generation: u64,
    reconcile_submission_on_snapshot: Option<u64>,
    approval_submitting: bool,
}

impl Default for StructuredState {
    fn default() -> Self {
        Self {
            rows: Vec::new(),
            reduced: ReducedState::default(),
            selected: None,
            expanded: HashSet::new(),
            scroll: 0,
            line_rows: Vec::new(),
            follow_selection: true,
            composer: TextInput::default(),
            error: None,
            submitting: false,
            submission_generation: 0,
            reconcile_submission_on_snapshot: None,
            approval_submitting: false,
        }
    }
}

impl StructuredState {
    pub fn apply(&mut self, event: StructuredEvent) {
        match event {
            StructuredEvent::TranscriptSnapshot { snapshot } => {
                let selected = self.selected.clone();
                self.rows = snapshot.rows;
                self.selected = selected
                    .filter(|id| self.rows.iter().any(|row| &row.id == id))
                    .or_else(|| {
                        self.rows
                            .iter()
                            .rev()
                            .find(|row| row.kind != TranscriptRowKind::Plan)
                            .map(|row| row.id.clone())
                    });
                self.expanded
                    .retain(|id| self.rows.iter().any(|row| &row.id == id));
            }
            StructuredEvent::ReducedState { state } => {
                let reconcile_submission = self
                    .reconcile_submission_on_snapshot
                    .take()
                    .is_some_and(|generation| {
                        self.submitting && generation == self.submission_generation
                    });
                self.apply_reduced(state, reconcile_submission);
            }
            StructuredEvent::ProjectionUpdate { transcript, state } => {
                if let Some(delta) = transcript {
                    self.apply_delta(delta);
                }
                self.apply_reduced(state, false);
            }
        }
    }

    fn apply_delta(&mut self, delta: TranscriptDelta) {
        match delta {
            TranscriptDelta::Append { row } => {
                let followed_tail = self.follow_selection
                    && self.selected.as_ref().is_some_and(|id| {
                        self.rows
                            .iter()
                            .rev()
                            .find(|old| old.kind != TranscriptRowKind::Plan)
                            .is_some_and(|old| &old.id == id)
                    });
                if let Some(index) = self.rows.iter().position(|old| old.id == row.id) {
                    self.rows[index] = row.clone();
                } else {
                    self.rows.push(row.clone());
                }
                if row.kind != TranscriptRowKind::Plan && (self.selected.is_none() || followed_tail)
                {
                    self.selected = Some(row.id);
                }
            }
            TranscriptDelta::Patch { row } => {
                if let Some(index) = self.rows.iter().position(|old| old.id == row.id) {
                    self.rows[index] = row.clone();
                } else {
                    self.rows.push(row.clone());
                }
            }
            TranscriptDelta::Remove { id } => {
                let removed = self.rows.iter().position(|row| row.id == id);
                self.rows.retain(|row| row.id != id);
                self.expanded.remove(&id);
                if self.selected.as_deref() == Some(&id) {
                    self.selected = removed.and_then(|index| {
                        let split = index.min(self.rows.len());
                        self.rows[split..]
                            .iter()
                            .find(|row| row.kind != TranscriptRowKind::Plan)
                            .or_else(|| {
                                self.rows[..split]
                                    .iter()
                                    .rev()
                                    .find(|row| row.kind != TranscriptRowKind::Plan)
                            })
                            .map(|row| row.id.clone())
                    });
                }
            }
        }
    }

    fn apply_reduced(&mut self, state: ReducedState, reconcile_submission: bool) {
        let previous_nonce = self.pending_nonce().map(str::to_owned);
        if reconcile_submission || state.turn.responding || state.turn.thinking {
            self.submitting = false;
        }
        self.reduced = state;
        let next_nonce = self.pending_nonce().map(str::to_owned);
        if next_nonce != previous_nonce {
            self.approval_submitting = false;
            if next_nonce.is_some() {
                self.selected = self
                    .rows
                    .iter()
                    .rev()
                    .find(|row| row.kind == TranscriptRowKind::Approval)
                    .map(|row| row.id.clone());
                self.follow_selection = true;
            }
        }
    }

    fn prepare_snapshot_reconciliation(&mut self) {
        self.reconcile_submission_on_snapshot =
            self.submitting.then_some(self.submission_generation);
    }

    pub fn select(&mut self, delta: isize) {
        let selectable: Vec<&TranscriptRow> = self
            .rows
            .iter()
            .filter(|row| row.kind != TranscriptRowKind::Plan)
            .collect();
        if selectable.is_empty() {
            self.selected = None;
            return;
        }
        let current = self
            .selected
            .as_ref()
            .and_then(|id| selectable.iter().position(|row| &row.id == id))
            .unwrap_or(selectable.len() - 1);
        let next = if delta >= 0 {
            current
                .saturating_add(delta as usize)
                .min(selectable.len() - 1)
        } else {
            current.saturating_sub(delta.unsigned_abs())
        };
        self.selected = Some(selectable[next].id.clone());
        self.follow_selection = true;
    }

    pub fn toggle_selected(&mut self) {
        let Some(id) = self.selected.clone() else {
            return;
        };
        let expandable = self.rows.iter().any(|row| {
            row.id == id
                && matches!(
                    row.kind,
                    TranscriptRowKind::ToolCall | TranscriptRowKind::ToolDiff
                )
        });
        if !expandable {
            return;
        }
        if !self.expanded.insert(id.clone()) {
            self.expanded.remove(&id);
        }
        self.follow_selection = true;
    }

    pub fn scroll_by(&mut self, lines: isize) {
        self.scroll = if lines >= 0 {
            self.scroll.saturating_sub(lines as usize)
        } else {
            self.scroll.saturating_add(lines.unsigned_abs())
        };
        self.follow_selection = false;
    }

    pub fn edit_composer(&mut self, key: &InputKey) {
        if !self.turn_active() {
            self.composer.handle_key(key);
            self.error = None;
        }
    }

    pub fn begin_prompt(&mut self) -> Option<String> {
        if self.turn_active() || self.composer.value().trim().is_empty() {
            return None;
        }
        let prompt = self.composer.value().to_string();
        self.composer = TextInput::default();
        self.submission_generation = self.submission_generation.wrapping_add(1);
        self.submitting = true;
        self.error = None;
        Some(prompt)
    }

    pub fn prompt_failed(&mut self, prompt: String, error: String) {
        if self.composer.value().is_empty() {
            self.composer = TextInput::with_value(prompt);
        }
        self.submitting = false;
        self.error = Some(error);
    }

    pub fn turn_active(&self) -> bool {
        self.submitting || self.reduced.turn.responding || self.reduced.turn.thinking
    }

    pub fn pending_nonce(&self) -> Option<&str> {
        self.reduced
            .pending_approval
            .as_ref()?
            .get("nonce")?
            .as_str()
    }

    pub fn approval_decision(&mut self, key: char) -> Option<(String, StructuredApprovalDecision)> {
        if self.approval_submitting {
            return None;
        }
        let selected = self.selected.as_deref()?;
        let pending_row = self
            .rows
            .iter()
            .rev()
            .find(|row| row.kind == TranscriptRowKind::Approval)?;
        if pending_row.id != selected {
            return None;
        }
        let decision = match key {
            'a' => StructuredApprovalDecision::AllowOnce,
            'A' => StructuredApprovalDecision::AllowForSession,
            'd' => StructuredApprovalDecision::Deny,
            _ => return None,
        };
        let offered = offered_decisions(&pending_row.data);
        if !offered.contains(&decision) {
            return None;
        }
        let nonce = self.pending_nonce()?.to_string();
        self.approval_submitting = true;
        self.error = None;
        Some((nonce, decision))
    }

    pub fn approval_failed(&mut self, nonce: &str, error: String) {
        if self.pending_nonce() != Some(nonce) {
            return;
        }
        self.approval_submitting = false;
        self.error = Some(error);
    }
}

pub fn offered_decisions(data: &serde_json::Value) -> Vec<StructuredApprovalDecision> {
    let mut out = Vec::new();
    for kind in data
        .get("options")
        .and_then(serde_json::Value::as_array)
        .into_iter()
        .flatten()
        .filter_map(|option| option.get("kind").and_then(serde_json::Value::as_str))
    {
        match kind {
            "allow_once" if !out.contains(&StructuredApprovalDecision::AllowOnce) => {
                out.push(StructuredApprovalDecision::AllowOnce)
            }
            "allow_always" if !out.contains(&StructuredApprovalDecision::AllowForSession) => {
                out.push(StructuredApprovalDecision::AllowForSession)
            }
            "reject_once" | "reject_always" if !out.contains(&StructuredApprovalDecision::Deny) => {
                out.push(StructuredApprovalDecision::Deny)
            }
            _ => {}
        }
    }
    out
}

/// Cheap shared handle for a structured projection connection.
#[derive(Clone)]
pub struct StructuredHandle {
    pub state: Arc<Mutex<StructuredState>>,
    connection: Arc<AtomicU8>,
    revision: Arc<AtomicU64>,
    stop: watch::Sender<bool>,
}

impl std::fmt::Debug for StructuredHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StructuredHandle")
            .field("connection", &self.connection())
            .field("revision", &self.revision())
            .finish_non_exhaustive()
    }
}

impl StructuredHandle {
    pub fn connection(&self) -> SessionState {
        match self.connection.load(Ordering::Acquire) {
            0 => SessionState::Connecting,
            1 => SessionState::Live,
            2 => SessionState::Reconnecting,
            _ => SessionState::Ended,
        }
    }

    pub fn revision(&self) -> u64 {
        self.revision.load(Ordering::Acquire)
    }

    pub fn end(&self) {
        self.connection
            .store(SessionState::Ended as u8, Ordering::Release);
        let _ = self.stop.send(true);
    }

    #[cfg(test)]
    pub(crate) fn test_handle(state: StructuredState) -> Self {
        let (stop, _) = watch::channel(false);
        Self {
            state: Arc::new(Mutex::new(state)),
            connection: Arc::new(AtomicU8::new(SessionState::Live as u8)),
            revision: Arc::new(AtomicU64::new(0)),
            stop,
        }
    }
}

#[derive(Debug)]
pub struct StructuredTask;

impl StructuredTask {
    pub fn spawn(ws_url: String) -> StructuredHandle {
        let state = Arc::new(Mutex::new(StructuredState::default()));
        let connection = Arc::new(AtomicU8::new(SessionState::Connecting as u8));
        let revision = Arc::new(AtomicU64::new(0));
        let (stop, mut stopped) = watch::channel(false);
        let handle = StructuredHandle {
            state: state.clone(),
            connection: connection.clone(),
            revision: revision.clone(),
            stop,
        };

        tokio::spawn(async move {
            let mut backoff = Duration::from_secs(1);
            loop {
                if *stopped.borrow() {
                    return;
                }
                state
                    .lock()
                    .expect("structured state mutex")
                    .prepare_snapshot_reconciliation();
                connection.store(SessionState::Connecting as u8, Ordering::Release);
                let connected = tokio::select! {
                    result = tokio_tungstenite::connect_async(&ws_url) => result,
                    _ = stopped.changed() => return,
                };
                let (mut ws, _) = match connected {
                    Ok(value) => value,
                    Err(error) => {
                        if terminal_handshake_error(&error) {
                            state.lock().expect("structured state mutex").error =
                                Some(format!("cannot open structured session: {error}"));
                            connection.store(SessionState::Ended as u8, Ordering::Release);
                            revision.fetch_add(1, Ordering::Release);
                            return;
                        }
                        connection.store(SessionState::Reconnecting as u8, Ordering::Release);
                        revision.fetch_add(1, Ordering::Release);
                        tokio::select! {
                            _ = tokio::time::sleep(super::session_task::jitter(backoff)) => {},
                            _ = stopped.changed() => return,
                        }
                        backoff = backoff.saturating_mul(2).min(Duration::from_secs(30));
                        continue;
                    }
                };
                backoff = Duration::from_secs(1);
                connection.store(SessionState::Live as u8, Ordering::Release);
                revision.fetch_add(1, Ordering::Release);

                loop {
                    let frame = tokio::select! {
                        frame = ws.next() => frame,
                        _ = stopped.changed() => return,
                    };
                    match frame {
                        Some(Ok(Message::Text(text))) => {
                            if let Ok(event) = serde_json::from_str::<StructuredEvent>(&text) {
                                state.lock().expect("structured state mutex").apply(event);
                                revision.fetch_add(1, Ordering::Release);
                            }
                        }
                        Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
                        _ => {}
                    }
                }
                connection.store(SessionState::Reconnecting as u8, Ordering::Release);
                revision.fetch_add(1, Ordering::Release);
                tokio::select! {
                    _ = tokio::time::sleep(super::session_task::jitter(backoff)) => {},
                    _ = stopped.changed() => return,
                }
                backoff = backoff.saturating_mul(2).min(Duration::from_secs(30));
            }
        });
        handle
    }
}

fn terminal_handshake_error(error: &WsError) -> bool {
    matches!(error, WsError::Http(response) if response.status().is_client_error())
}

#[cfg(test)]
mod tests {
    use super::*;
    use maestro_protocol::{TranscriptSnapshot, TurnState};
    use serde_json::json;

    fn row(id: &str, kind: TranscriptRowKind, data: serde_json::Value) -> TranscriptRow {
        TranscriptRow {
            id: id.into(),
            kind,
            data,
        }
    }

    #[test]
    fn snapshots_replace_and_deltas_address_rows_by_id() {
        let mut state = StructuredState::default();
        state.apply(StructuredEvent::TranscriptSnapshot {
            snapshot: TranscriptSnapshot {
                cursor: 2,
                rows: vec![row("a", TranscriptRowKind::AgentText, json!({}))],
            },
        });
        state.apply(StructuredEvent::ProjectionUpdate {
            transcript: Some(TranscriptDelta::Patch {
                row: row("a", TranscriptRowKind::AgentText, json!({"n": 2})),
            }),
            state: ReducedState::default(),
        });
        state.apply(StructuredEvent::ProjectionUpdate {
            transcript: Some(TranscriptDelta::Append {
                row: row("b", TranscriptRowKind::Notice, json!({})),
            }),
            state: ReducedState::default(),
        });
        assert_eq!(state.rows.len(), 2);
        assert_eq!(state.rows[0].data["n"], 2);

        state.apply(StructuredEvent::TranscriptSnapshot {
            snapshot: TranscriptSnapshot {
                cursor: 5,
                rows: vec![row("b", TranscriptRowKind::Notice, json!({"fresh": true}))],
            },
        });
        assert_eq!(state.rows.len(), 1);
        assert_eq!(state.rows[0].data["fresh"], true);
    }

    #[test]
    fn selection_survives_patch_and_moves_after_remove() {
        let mut state = StructuredState {
            rows: vec![
                row("a", TranscriptRowKind::Notice, json!({})),
                row("b", TranscriptRowKind::ToolCall, json!({})),
                row("plan", TranscriptRowKind::Plan, json!({})),
            ],
            selected: Some("a".into()),
            ..StructuredState::default()
        };
        state.select(1);
        assert_eq!(state.selected.as_deref(), Some("b"));
        state.toggle_selected();
        assert!(state.expanded.contains("b"));
        state.apply_delta(TranscriptDelta::Remove { id: "b".into() });
        assert_eq!(state.selected.as_deref(), Some("a"));

        state.apply_delta(TranscriptDelta::Remove { id: "a".into() });
        assert_eq!(
            state.selected, None,
            "collapsed plan rows are not selectable"
        );
    }

    #[test]
    fn approval_keys_only_offer_server_backed_decisions() {
        let mut state = StructuredState::default();
        state.rows.push(row(
            "approval:1",
            TranscriptRowKind::Approval,
            json!({"options":[
                {"kind":"allow_once"}, {"kind":"reject_always"}
            ]}),
        ));
        state.selected = Some("approval:1".into());
        state.reduced.pending_approval = Some(json!({"nonce":"secret"}));
        assert_eq!(
            state.approval_decision('A'),
            None,
            "allow-for-session was not offered"
        );
        assert_eq!(
            state.approval_decision('a'),
            Some(("secret".into(), StructuredApprovalDecision::AllowOnce))
        );
    }

    #[test]
    fn resolved_approval_ignores_its_late_failure() {
        let mut state = StructuredState::default();
        state.rows.push(row(
            "approval:1",
            TranscriptRowKind::Approval,
            json!({"options":[{"kind":"allow_once"}]}),
        ));
        state.selected = Some("approval:1".into());
        state.reduced.pending_approval = Some(json!({"nonce":"old"}));
        assert!(state.approval_decision('a').is_some());

        state.apply(StructuredEvent::ProjectionUpdate {
            transcript: None,
            state: ReducedState::default(),
        });
        state.approval_failed("old", "response lost".into());

        assert_eq!(state.error, None);
        assert!(!state.approval_submitting);
    }

    #[test]
    fn superseded_approval_ignores_its_late_failure() {
        let mut state = StructuredState::default();
        state.rows.push(row(
            "approval:1",
            TranscriptRowKind::Approval,
            json!({"options":[{"kind":"allow_once"}]}),
        ));
        state.selected = Some("approval:1".into());
        state.reduced.pending_approval = Some(json!({"nonce":"old"}));
        assert!(state.approval_decision('a').is_some());

        state.rows.push(row(
            "approval:2",
            TranscriptRowKind::Approval,
            json!({"options":[{"kind":"allow_once"}]}),
        ));
        state.apply(StructuredEvent::ProjectionUpdate {
            transcript: None,
            state: ReducedState {
                pending_approval: Some(json!({"nonce":"new"})),
                ..ReducedState::default()
            },
        });
        assert_eq!(
            state.approval_decision('a'),
            Some(("new".into(), StructuredApprovalDecision::AllowOnce))
        );

        state.approval_failed("old", "response lost".into());

        assert_eq!(state.error, None);
        assert!(state.approval_submitting);
    }

    #[test]
    fn matching_approval_surfaces_its_failure() {
        let mut state = StructuredState::default();
        state.rows.push(row(
            "approval:1",
            TranscriptRowKind::Approval,
            json!({"options":[{"kind":"allow_once"}]}),
        ));
        state.selected = Some("approval:1".into());
        state.reduced.pending_approval = Some(json!({"nonce":"current"}));
        assert!(state.approval_decision('a').is_some());

        state.approval_failed("current", "request failed".into());

        assert_eq!(state.error.as_deref(), Some("request failed"));
        assert!(!state.approval_submitting);
    }

    #[test]
    fn composer_is_gated_by_the_projected_turn_lifecycle() {
        let mut state = StructuredState::default();
        state.edit_composer(&InputKey::char('h'));
        state.edit_composer(&InputKey::char('i'));
        assert_eq!(state.begin_prompt().as_deref(), Some("hi"));
        assert!(state.turn_active());
        state.prompt_failed("hi".into(), "nope".into());
        assert_eq!(state.composer.value(), "hi");

        state.reduced.turn = TurnState {
            responding: true,
            ..TurnState::default()
        };
        state.edit_composer(&InputKey::char('x'));
        assert_eq!(state.composer.value(), "hi");
    }

    #[test]
    fn authoritative_idle_snapshot_reconciles_a_wholly_missed_turn() {
        let mut state = StructuredState::default();
        state.edit_composer(&InputKey::char('h'));
        state.edit_composer(&InputKey::char('i'));
        assert_eq!(state.begin_prompt().as_deref(), Some("hi"));

        state.apply(StructuredEvent::ProjectionUpdate {
            transcript: None,
            state: ReducedState::default(),
        });
        assert!(
            state.turn_active(),
            "an ordinary idle update must not reopen the composer before turn begin"
        );

        state.prepare_snapshot_reconciliation();
        state.apply(StructuredEvent::TranscriptSnapshot {
            snapshot: TranscriptSnapshot {
                cursor: 0,
                rows: Vec::new(),
            },
        });
        state.apply(StructuredEvent::ReducedState {
            state: ReducedState::default(),
        });
        assert!(
            !state.turn_active(),
            "a reconnect snapshot must settle a submission whose whole turn was missed"
        );
    }

    #[test]
    fn stale_handshake_state_does_not_settle_a_new_submission() {
        let mut state = StructuredState::default();
        state.prepare_snapshot_reconciliation();
        state.apply(StructuredEvent::TranscriptSnapshot {
            snapshot: TranscriptSnapshot {
                cursor: 0,
                rows: Vec::new(),
            },
        });

        state.edit_composer(&InputKey::char('h'));
        state.edit_composer(&InputKey::char('i'));
        assert_eq!(state.begin_prompt().as_deref(), Some("hi"));
        state.apply(StructuredEvent::ReducedState {
            state: ReducedState::default(),
        });

        assert!(
            state.turn_active(),
            "state captured before this submission must not reopen the composer"
        );
    }

    #[test]
    fn failed_submission_reconciliation_does_not_settle_its_successor() {
        let mut state = StructuredState::default();
        state.edit_composer(&InputKey::char('a'));
        assert_eq!(state.begin_prompt().as_deref(), Some("a"));
        state.prepare_snapshot_reconciliation();

        state.prompt_failed("a".into(), "nope".into());
        state.composer = TextInput::with_value("b");
        assert_eq!(state.begin_prompt().as_deref(), Some("b"));

        state.apply(StructuredEvent::TranscriptSnapshot {
            snapshot: TranscriptSnapshot {
                cursor: 0,
                rows: Vec::new(),
            },
        });
        state.apply(StructuredEvent::ReducedState {
            state: ReducedState::default(),
        });

        assert!(
            state.turn_active(),
            "a reconnect for the failed submission must not settle its successor"
        );
    }
}