ai-app/iris/android-app/src/transcript_client.rs

//! RUST.md's I5 Android integration: `transcript-ui`'s screen filling the
//! whole window on android-view, against a real `ai-server` through
//! `client-core` -- the missing half `iris-android-app` (I2) only had for
//! `tabs-ui` until now. Behind the `transcript-screen` Cargo feature so the
//! plain build (`cargo ndk build`, no `--features`) stays exactly the tabs
//! demo I2/I4 already measured against.
//!
//! **Deliberate simplification, recorded rather than left to be
//! rediscovered (RUST.md's I5 box has the full account)**: there is no
//! session list and no enrollment UI here. The server, port, token and
//! pinned CA are baked in at build time (`build.rs`'s
//! `AI_APP_TRANSCRIPT_HOST`/`_PORT`/`_TOKEN`/`AI_APP_CA`), and the first
//! session `ApiClient::fetch_sessions` returns is opened automatically --
//! there is nothing to tap to get there, which is what `transcript-bench.sh`
//! and `ui-trace` need to land straight on the screen under test. A real
//! app needs `desktop-app`'s `EnrolledServer`/QR-link flow or E3's
//! Keystore-sealed `ServerConfig.kt`; building a second one of those was
//! not this pass's job.
//!
//! **Reuses `iris/desktop-app`'s `app.rs` shape almost exactly** --
//! `fold_event`/`group_tool_runs`/`fold_page`/`raw_seq` from
//! `client_core::transcript_fold`, a `generation` counter guarding against
//! a stale background response. What differs is only the redraw
//! mechanism: android-view has no `winit::EventLoopProxy`, so this uses
//! `iris::task::Tasks::redraw_handle` (new, added alongside this box) to
//! request a frame after each `TaskCtx::update` instead of relying on
//! `Tasks::spawn`'s single end-of-future redraw -- see that method's own
//! doc for why.
//!
//! **Streaming no longer costs a full rebuild** (fixed after the P0 gate
//! showed why it mattered -- 20 events/second means 20 rebuilds/second of
//! a ~3,200-row transcript otherwise): `apply_event` calls
//! `transcript_ui::TranscriptScreen::apply` with the item list before and
//! after `fold_event`, which updates only the row(s) that actually
//! changed (almost always just the one open assistant message) instead of
//! refolding and rebuilding every row. `rebuild_transcript` still runs
//! the whole widget tree once, for the opening page and for `apply`'s own
//! rare regroup fallback.

use client_core::api::{ApiClient, UreqTransport};
use client_core::event_stream::{StreamItem, follow_session_events};
use client_core::transcript_fold::{TranscriptItem, fold_event, fold_page, group_tool_runs};
use event_model::SeqEvent;
use iris::android::{AndroidAppState, AndroidRsc, AndroidUiState, HasAndroidUiState};
use iris::prelude::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

mod pinned {
    include!(concat!(env!("OUT_DIR"), "/pinned_config.rs"));
}

pub struct TranscriptClient {
    ui_state: AndroidUiState,
    /// The screen's own content -- everything under the fixed
    /// [`frame_report_controls`] bar, which is built once (`new`, below)
    /// and never touched by `show_message`/`rebuild_transcript`'s own
    /// `set` calls the way `desktop-app`'s `transcript_ptr` isn't touched
    /// by rebuilding the session list beside it.
    content: WeakWidget<WidgetPtr>,
    screen: Option<transcript_ui::TranscriptScreen>,
    /// The folded transcript as of the last rebuild -- kept here (not
    /// re-derived) for the same reason `desktop-app`'s `Client::items`
    /// exists: a live `StreamEvent` only carries one new wire event, and
    /// `fold_event` needs everything folded so far to fold it in.
    items: Vec<TranscriptItem>,
    /// The session currently open -- `None` only before the first fetch
    /// resolves. Read back by `apply_event`'s rebuild, which has no session
    /// id of its own (a live `SeqEvent` doesn't carry one).
    session_id: Option<String>,
    /// Bumped every time a new session load starts; a background response
    /// checks it before touching state, so a slow reply for a session this
    /// screen has moved on from can't overwrite what replaced it. There is
    /// only ever one session here (no list to switch away to), but the
    /// guard still matters for the *first* fetch racing a `stop`/`start`.
    generation: Arc<AtomicU64>,
}

impl HasAndroidUiState for TranscriptClient {
    fn android_state(&self) -> &AndroidUiState {
        &self.ui_state
    }
    fn android_state_mut(&mut self) -> &mut AndroidUiState {
        &mut self.ui_state
    }
}

/// Builds one `UreqTransport` from the config `build.rs` baked in. Called
/// twice per session load, same as `desktop-app`'s `build_transport`
/// closure -- `ApiClient` and the live-stream follow each need their own,
/// since `UreqTransport` holds its own `ureq::Agent`.
fn build_transport() -> Result<UreqTransport, String> {
    let base_url = format!("https://{}:{}", pinned::HOST, pinned::PORT);
    UreqTransport::new(
        base_url,
        pinned::TOKEN.to_string(),
        pinned::CA_PEM.as_bytes(),
    )
    .map_err(|e| e.to_string())
}

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

/// The two named controls RUST.md's I5 box ("Measurements taken" (b))
/// drives by name over `ui-trace`, e.g. `ui-trace record --do "tap 'Frame
/// report'"`. `dumpsys gfxinfo` cannot see this screen's own GPU-drawn
/// frames at all -- this is the screen's own equivalent of the Compose
/// app's "Copy render timings" control, logged rather than clipboarded
/// (no clipboard wiring exists here) under this crate's own fixed
/// `android_logger` tag (`iris-android-app`, `lib.rs`'s `JNI_OnLoad`),
/// grep-able on the fixed string `"iris frame report"` the way
/// `transcript-bench.sh` greps `"ai-app render report"`.
fn frame_report_controls(rsc: &mut AndroidRsc<TranscriptClient>) -> WeakWidget {
    type Rsc = AndroidRsc<TranscriptClient>;
    let report_rect = rect(Color::rgb(50, 50, 60))
        .on(
            CursorSense::click(),
            |ctx: EventIdCtx<'_, Rsc, _, _>, _rsc: &mut Rsc| match ctx
                .state
                .android_state()
                .frame_report
                .report()
            {
                Some(stats) => log::info!("iris frame report: {stats}"),
                None => log::info!(
                    "iris frame report: no frames recorded -- scroll first, then press this"
                ),
            },
        )
        .label("Frame report");
    let report = (
        report_rect,
        wtext("Frame report").size(18).text_align(Align::CENTER),
    )
        .stack()
        .pad(8)
        .add(rsc);

    let reset_rect = rect(Color::rgb(70, 40, 40))
        .on(
            CursorSense::click(),
            |ctx: EventIdCtx<'_, Rsc, _, _>, _rsc: &mut Rsc| {
                ctx.state.android_state_mut().frame_report.reset();
                log::info!("iris frame report: reset");
            },
        )
        .label("Reset frame report");
    let reset = (
        reset_rect,
        wtext("Reset").size(18).text_align(Align::CENTER),
    )
        .stack()
        .pad(8)
        .add(rsc);

    (report, reset).span(Dir::RIGHT).height(56).add(rsc)
}

impl AndroidAppState for TranscriptClient {
    fn new(mut ui_state: AndroidUiState, rsc: &mut AndroidRsc<Self>) -> Self {
        let content = WidgetPtr::new().add(rsc);
        let loading = placeholder(rsc, "Loading sessions...");
        content(rsc).set(loading);

        let tree = (frame_report_controls(rsc), content.height(rest(1)))
            .span(Dir::DOWN)
            .add_strong(rsc)
            .any();
        ui_state.set_root(tree);

        let mut client = Self {
            ui_state,
            content,
            screen: None,
            items: Vec::new(),
            session_id: None,
            generation: Arc::new(AtomicU64::new(0)),
        };
        client.spawn_fetch_sessions(rsc);
        client
    }

    fn back_pressed(&mut self, _rsc: &mut AndroidRsc<Self>, _render: &mut UiRenderState) -> bool {
        // No screen stack of its own -- same "let the activity finish"
        // answer `iris-android-app`'s tabs `Client` already gives.
        false
    }
}

impl TranscriptClient {
    fn show_message(&mut self, rsc: &mut AndroidRsc<Self>, message: &str) {
        let widget = placeholder(rsc, message);
        (self.content)(rsc).set(widget);
        self.screen = None;
    }

    fn spawn_fetch_sessions(&mut self, rsc: &mut AndroidRsc<Self>) {
        let redraw = rsc.tasks.redraw_handle();
        let my_generation = self.generation.load(Ordering::SeqCst);
        let generation = self.generation.clone();
        rsc.spawn_task(async move |mut ctx| {
            let outcome = match build_transport() {
                Ok(transport) => ApiClient::new(transport)
                    .fetch_sessions()
                    .map_err(|e| e.to_string()),
                Err(e) => Err(format!("couldn't set up TLS: {e}")),
            };
            ctx.update(move |state: &mut TranscriptClient, rsc| {
                if generation.load(Ordering::SeqCst) != my_generation {
                    return;
                }
                match outcome {
                    Ok(sessions) => match sessions.into_iter().next() {
                        Some(session) => state.select_session(rsc, session.id),
                        None => state.show_message(rsc, "No sessions on the sandbox server."),
                    },
                    Err(message) => {
                        state.show_message(rsc, &format!("Couldn't list sessions: {message}"))
                    }
                }
            });
            redraw.request_redraw();
        });
    }

    /// Loads the opening page, then follows the live SSE stream for the
    /// rest of this session's life -- `desktop-app`'s `select_session`
    /// almost verbatim, with `Proxy::send_event` replaced by `ctx.update` +
    /// `redraw.request_redraw()` (see this module's doc).
    fn select_session(&mut self, rsc: &mut AndroidRsc<Self>, session_id: String) {
        let my_generation = self.generation.fetch_add(1, Ordering::SeqCst) + 1;
        self.items.clear();
        self.session_id = Some(session_id.clone());
        self.show_message(rsc, "Loading transcript...");

        let redraw = rsc.tasks.redraw_handle();
        let live_generation = self.generation.clone();
        rsc.spawn_task(async move |mut ctx| {
            let transports =
                build_transport().and_then(|rest| build_transport().map(|stream| (rest, stream)));
            let (rest, stream_transport) = match transports {
                Ok(pair) => pair,
                Err(e) => {
                    let message = format!("couldn't set up TLS: {e}");
                    ctx.update(move |state: &mut TranscriptClient, rsc| {
                        if live_generation.load(Ordering::SeqCst) == my_generation {
                            state.show_message(rsc, &message);
                        }
                    });
                    redraw.request_redraw();
                    return;
                }
            };
            let api = ApiClient::new(rest);

            // The most recent 200 events, coalesced -- the same page size
            // `desktop-app` uses; RUST.md's I3/history-paging work is what
            // a real scrollback would reuse (out of scope here, same as
            // E4).
            let page: Result<Vec<serde_json::Value>, String> = api
                .fetch_transcript_page(&session_id, None, 200, true)
                .map_err(|e| e.to_string());
            // The wire `seq` of the last line, not a folded item's `seq()`
            // -- see `client_core::transcript_fold::raw_seq`'s doc for why
            // resuming from the latter re-delivers deltas already folded
            // into an in-progress reply.
            let after = page
                .as_ref()
                .ok()
                .and_then(|values| values.last())
                .and_then(client_core::transcript_fold::raw_seq)
                .unwrap_or(0);
            let result = page.and_then(|values| fold_page(&values));

            {
                let live_generation = live_generation.clone();
                ctx.update(move |state: &mut TranscriptClient, rsc| {
                    if live_generation.load(Ordering::SeqCst) != my_generation {
                        return;
                    }
                    match result {
                        Ok(items) => {
                            state.items = items;
                            state.rebuild_transcript(rsc);
                        }
                        Err(message) => {
                            state.show_message(rsc, &format!("Couldn't load transcript: {message}"))
                        }
                    }
                });
            }
            redraw.request_redraw();

            if live_generation.load(Ordering::SeqCst) != my_generation {
                return;
            }
            // The outer closure here is an `FnMut` -- `follow_session_events`
            // calls it once per line -- so it captures `live_generation` by
            // move and re-clones it for each inner `ctx.update` closure
            // rather than moving a shared `stop`-style helper into itself:
            // a value moved out of an `FnMut`'s captures on one call leaves
            // nothing there for the next.
            let _ =
                follow_session_events(
                    &stream_transport,
                    &session_id,
                    after,
                    move |item| match item {
                        StreamItem::Open | StreamItem::Reset => {
                            live_generation.load(Ordering::SeqCst) == my_generation
                        }
                        StreamItem::Event { event, .. } => {
                            if live_generation.load(Ordering::SeqCst) != my_generation {
                                return false;
                            }
                            let live_generation = live_generation.clone();
                            ctx.update(move |state: &mut TranscriptClient, rsc| {
                                if live_generation.load(Ordering::SeqCst) != my_generation {
                                    return;
                                }
                                state.apply_event(rsc, &event);
                            });
                            redraw.request_redraw();
                            true
                        }
                    },
                );
        });
    }

    /// Rebuilds the whole widget tree from `self.items` -- same tradeoff as
    /// `desktop-app`'s `rebuild_transcript` (this module's doc comment).
    /// Reads `self.session_id` rather than taking one, since every caller
    /// (the opening page, and every live event) already has it set there.
    fn rebuild_transcript(&mut self, rsc: &mut AndroidRsc<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) = transcript_ui::build_tree(rsc, rows);

        if let Some(text) = in_progress {
            screen.composer.field.edit(rsc).set(&text);
        }
        if let Some(session_id) = self.session_id.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.content)(rsc).set(tree);
        self.screen = Some(screen);
    }

    fn apply_event(&mut self, rsc: &mut AndroidRsc<Self>, event: &SeqEvent) {
        let old_items = self.items.clone();
        self.items = fold_event(&self.items, event);
        match &self.screen {
            // The common path: update only the row(s) that actually
            // changed instead of refolding and rebuilding all ~3,200 of
            // them per event (RUST.md's P0 streaming-phase fix).
            Some(screen) => screen.apply(rsc, &old_items, &self.items),
            // No screen yet (the opening page hasn't landed) -- build one
            // the ordinary way once it has.
            None => self.rebuild_transcript(rsc),
        }
    }

    fn send_message(&mut self, session_id: String, text: String) {
        std::thread::spawn(move || {
            if let Ok(transport) = build_transport() {
                let api = ApiClient::new(transport);
                let _ = api.send_message(&session_id, &text, &[]);
            }
        });
    }
}