ai-app/app/src/desktop/app.rs

use crate::client::api::{ApiClient, SessionSummary, UreqTransport};
use crate::client::event_stream::{StreamItem, follow_session_events};
use crate::client::transcript_fold::{
    TranscriptItem, fold_event, fold_page, group_tool_runs, raw_seq,
};
use event_model::SeqEvent;
use iris::prelude::*;
use std::{
    process,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    thread,
};

const LIST_WIDTH: f32 = 260.0;

enum AppEvent {
    Sessions(Result<Vec<SessionSummary>, String>),
    TranscriptLoaded {
        session_id: String,
        generation: u64,
        result: Result<Vec<TranscriptItem>, String>,
    },
    StreamEvent {
        session_id: String,
        generation: u64,
        event: SeqEvent,
    },
    StreamEnded {
        session_id: String,
        generation: u64,
        message: Option<String>,
    },
    SendFailed(String),
}

pub fn run() {
    DesktopApp::<Client>::run();
}

#[derive(DesktopUiState)]
struct Client {
    ui_state: DesktopUiState,
    api: Arc<ApiClient<UreqTransport>>,
    stream_transport: Arc<UreqTransport>,
    updates: TaskCtx<StdRsc<Self>>,
    sessions: Vec<SessionSummary>,
    selected: Option<String>,
    items: Vec<TranscriptItem>,
    list_ptr: WeakWidget<WidgetPtr>,
    transcript_ptr: WeakWidget<WidgetPtr>,
    screen: Option<crate::ui::TranscriptScreen>,
    generation: Arc<AtomicU64>,
}

impl DesktopAppState for Client {
    fn new(mut ui_state: DesktopUiState, rsc: &mut StdRsc<Self>) -> Self {
        let (server, ca_pem) = super::startup::load_startup_config().unwrap_or_else(|e| {
            eprintln!("desktop-app: {e}");
            process::exit(2);
        });
        let build_transport =
            || UreqTransport::new(server.base_url(), server.token.clone(), &ca_pem);
        let (rest_transport, stream_transport) = build_transport()
            .and_then(|rest| build_transport().map(|stream| (rest, stream)))
            .unwrap_or_else(|e| {
                eprintln!(
                    "desktop-app: couldn't set up TLS to {}: {e}",
                    server.base_url()
                );
                process::exit(1);
            });
        let api = Arc::new(ApiClient::new(rest_transport));
        let stream_transport = Arc::new(stream_transport);

        let list_ptr = WidgetPtr::new().add(rsc);
        let transcript_ptr = WidgetPtr::new().add(rsc);
        let loading = placeholder(rsc, "Loading sessions...");
        transcript_ptr(rsc).set(loading);

        (list_ptr.width(LIST_WIDTH), transcript_ptr.width(rest(1)))
            .span(Dir::RIGHT)
            .set_root(rsc, &mut ui_state);

        let client = Self {
            ui_state,
            api,
            stream_transport,
            updates: rsc.tasks.context(),
            sessions: Vec::new(),
            selected: None,
            items: Vec::new(),
            list_ptr,
            transcript_ptr,
            screen: None,
            generation: Arc::new(AtomicU64::new(0)),
        };
        client.spawn_fetch_sessions();
        client
    }
}

impl Client {
    fn apply_event(&mut self, event: AppEvent, rsc: &mut StdRsc<Self>) {
        match event {
            AppEvent::Sessions(Ok(sessions)) => {
                self.sessions = sessions;
                self.rebuild_list(rsc);
                if self.selected.is_none() {
                    self.show_message(rsc, "Select a session.");
                }
            }
            AppEvent::Sessions(Err(message)) => {
                self.show_message(rsc, &format!("Couldn't list sessions: {message}"));
            }
            AppEvent::TranscriptLoaded {
                session_id,
                generation,
                result,
            } => {
                if self.current(&session_id, generation) {
                    match result {
                        Ok(items) => {
                            self.items = items;
                            self.rebuild_transcript(rsc);
                        }
                        Err(message) => {
                            self.show_message(
                                rsc,
                                &format!("Couldn't load {session_id}: {message}"),
                            );
                        }
                    }
                }
            }
            AppEvent::StreamEvent {
                session_id,
                generation,
                event,
            } => {
                if self.current(&session_id, generation) {
                    let old_items = self.items.clone();
                    self.items = fold_event(&self.items, &event);
                    match self.screen.as_mut() {
                        Some(screen) => screen.apply(rsc, &old_items, &self.items),
                        None => self.rebuild_transcript(rsc),
                    }
                }
            }
            AppEvent::StreamEnded {
                session_id,
                generation,
                message: Some(message),
            } => {
                if self.current(&session_id, generation) {
                    eprintln!("desktop-app: {session_id}'s live connection ended: {message}");
                }
            }
            AppEvent::StreamEnded { .. } => {}
            AppEvent::SendFailed(message) => {
                eprintln!("desktop-app: couldn't send: {message}");
            }
        }
    }

    fn current(&self, session_id: &str, generation: u64) -> bool {
        self.selected.as_deref() == Some(session_id)
            && self.generation.load(Ordering::SeqCst) == generation
    }

    fn show_message(&mut self, rsc: &mut StdRsc<Self>, message: &str) {
        let widget = placeholder(rsc, message);
        (self.transcript_ptr)(rsc).set(widget);
    }

    fn spawn_fetch_sessions(&self) {
        let api = self.api.clone();
        let mut updates = self.updates.clone();
        thread::spawn(move || {
            let result = api.fetch_sessions().map_err(|e| e.to_string());
            updates.update(move |state: &mut Client, rsc| {
                state.apply_event(AppEvent::Sessions(result), rsc);
            });
        });
    }

    fn rebuild_list(&mut self, rsc: &mut StdRsc<Self>) {
        let list = Span::empty(Dir::DOWN).gap(2).add(rsc);
        for session in &self.sessions {
            let selected = self.selected.as_deref() == Some(session.id.as_str());
            let row = session_row(rsc, session, selected);
            list(rsc).push(row);
        }
        let tree = list
            .background(rect(Srgba8::rgb(24, 24, 28)))
            .add_strong(rsc)
            .any();
        (self.list_ptr)(rsc).set(tree);
    }

    fn select_session(&mut self, rsc: &mut StdRsc<Self>, session_id: String) {
        let generation = self.generation.fetch_add(1, Ordering::SeqCst) + 1;
        self.selected = Some(session_id.clone());
        self.items.clear();
        self.screen = None;
        self.rebuild_list(rsc);
        self.show_message(rsc, "Loading transcript...");

        let api = self.api.clone();
        let stream_transport = self.stream_transport.clone();
        let mut updates = self.updates.clone();
        let live_generation = self.generation.clone();
        thread::spawn(move || {
            let page: Result<Vec<serde_json::Value>, String> = api
                .fetch_transcript_page(&session_id, None, 200, true)
                .map_err(|e| e.to_string());
            let after = page
                .as_ref()
                .ok()
                .and_then(|values| raw_seq(values.last()?))
                .unwrap_or(0);
            let result = page.and_then(|values| fold_page(&values));
            let loaded_session_id = session_id.clone();
            updates.update(move |state: &mut Client, rsc| {
                state.apply_event(
                    AppEvent::TranscriptLoaded {
                        session_id: loaded_session_id,
                        generation,
                        result,
                    },
                    rsc,
                );
            });

            let stop = || live_generation.load(Ordering::SeqCst) != generation;
            if stop() {
                return;
            }
            let outcome =
                follow_session_events(&*stream_transport, &session_id, after, |item| match item {
                    StreamItem::Open | StreamItem::Reset => !stop(),
                    StreamItem::Event { event, .. } => {
                        if stop() {
                            return false;
                        }
                        let event_session_id = session_id.clone();
                        updates.update(move |state: &mut Client, rsc| {
                            state.apply_event(
                                AppEvent::StreamEvent {
                                    session_id: event_session_id,
                                    generation,
                                    event,
                                },
                                rsc,
                            );
                        });
                        true
                    }
                });
            updates.update(move |state: &mut Client, rsc| {
                state.apply_event(
                    AppEvent::StreamEnded {
                        session_id,
                        generation,
                        message: outcome.err().map(|e| e.to_string()),
                    },
                    rsc,
                );
            });
        });
    }

    fn send_message(&mut self, session_id: String, text: String) {
        let api = self.api.clone();
        let mut updates = self.updates.clone();
        thread::spawn(move || {
            if let Err(e) = api.send_message(&session_id, &text, &[]) {
                let message = e.to_string();
                updates.update(move |state: &mut Client, rsc| {
                    state.apply_event(AppEvent::SendFailed(message), rsc);
                });
            }
        });
    }

    fn rebuild_transcript(&mut self, rsc: &mut StdRsc<Self>) {
        let in_progress = self
            .screen
            .as_ref()
            .map(|screen| screen.composer.field.edit(rsc).text.text().to_string())
            .filter(|t| !t.is_empty());

        let rows = group_tool_runs(&self.items);
        let (screen, tree) = crate::ui::build_tree(rsc, rows);

        if let Some(text) = in_progress {
            screen.composer.field.edit(rsc).set(&text);
        }
        if let Some(session_id) = self.selected.clone() {
            let field = screen.composer.field;
            rsc.register_event(field, Submit, move |ctx, rsc| {
                let text = field.edit(rsc).take();
                let text = text.trim().to_string();
                if !text.is_empty() {
                    ctx.state.send_message(session_id.clone(), text);
                }
            });
        }

        (self.transcript_ptr)(rsc).set(tree);
        self.screen = Some(screen);
    }
}

fn session_row(rsc: &mut StdRsc<Client>, session: &SessionSummary, selected: bool) -> StrongWidget {
    let bg = if selected {
        Srgba8::rgb(58, 90, 138)
    } else {
        Srgba8::rgb(38, 38, 44)
    };
    let id = session.id.clone();
    let label = format!("{}\n{}", session.title, session.status);
    wtext(label)
        .color(PaintId::WHITE)
        .wrap(true)
        .pad(10)
        .width(rest(1))
        .background(rect(bg))
        .on(
            CursorSense::click(),
            move |ctx, rsc: &mut StdRsc<Client>| {
                ctx.state.select_session(rsc, id.clone());
            },
        )
        .add_strong(rsc)
        .any()
}

fn placeholder(rsc: &mut StdRsc<Client>, message: &str) -> StrongWidget {
    wtext(message.to_string())
        .color(PaintId::WHITE)
        .wrap(true)
        .pad(16)
        .add_strong(rsc)
        .any()
}