cydonia 0.1.2

//! One agent session bridged into the UI.
//!
//! The connection opens on the ACP runtime and the `Session` it yields is
//! held here, so dropping `ChatSession` tears the connection (and the agent
//! process) down. A foreground pump drains the ACP event channel in coalesced
//! batches with a 120ms frame floor while streaming — one notify per frame,
//! not per chunk.
//!
//! A session outlives its connection. [`ChatSession::flush`] writes the
//! transcript whenever a turn settles, and the agent's own session id goes
//! with it, so a relaunch reads the session back and `session/load` can pick
//! the conversation up where it stopped.

use crate::{
    agent::acp::{self, Event, Launch, Reply, Session},
    model::{settings, workspace::Workspace},
    view::component::transcript,
};
use anyhow::anyhow;
use artifact::{
    project::{Project as _, fs},
    session::{
        chat::{ChatItem, PlanStatus, ToolStatus},
        record::Record,
    },
};
use bezel::gpui::{Context, Task};
use cacp::schema::{
    ContentBlock, MaybeUndefined, PermissionOptionKind, PlanEntryStatus, RequestPermissionRequest,
    RequestPermissionResponse, SessionConfigKind, SessionConfigOption, SessionConfigOptionValue,
    SessionModeState, SessionUpdate, StopReason, ToolCallContent, ToolCallStatus,
};
use std::{
    collections::VecDeque,
    path::PathBuf,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

const STREAM_FRAME: Duration = Duration::from_millis(120);

/// How much of the context window the conversation has taken, as the agent
/// counts it.
///
/// Runtime only, and not written to the record: the file carries the
/// transcript, and whichever agent reads it back counts for itself.
#[derive(Clone, Copy)]
pub struct Usage {
    pub used: u64,
    pub size: u64,
}

impl Usage {
    /// How full, as a fraction. A window of no size is an agent that does not
    /// count, and there is nothing to draw for it.
    pub fn fraction(&self) -> Option<f32> {
        (self.size > 0).then(|| (self.used as f32 / self.size as f32).clamp(0., 1.))
    }
}

/// The turn in flight: when it started, and what the context had spent by then.
///
/// Runtime only, like [`Usage`], and for the same reason — the two numbers it
/// exists to count from are only read while a turn is running, and a relaunch
/// has none.
#[derive(Clone, Copy)]
struct Flight {
    at: SystemTime,
    used: u64,
}

/// One way to answer a permission request. `kind` is what decides how the
/// button paints — allow and reject must not look alike.
pub struct Choice {
    pub id: String,
    pub name: String,
    pub kind: PermissionOptionKind,
}

/// Whether the session can talk to an agent right now.
pub enum Connection {
    /// No agent process: read back from disk, archived, or given up on.
    Idle,
    Connecting,
    Live(Box<Session>),
    /// The agent went away on its own — a crash, or a closed stdout.
    Lost,
}

/// A slash command the agent offers for this session — what to type, and what
/// it does. The description is required by the protocol, so a picker can
/// always say what a name does not.
#[derive(Clone, PartialEq, Eq)]
pub struct Command {
    pub name: String,
    pub description: String,
}

pub struct PermissionPrompt {
    pub title: String,
    pub options: Vec<Choice>,
    /// Whether the answer should stand for every call like this one, rather
    /// than for this one alone — the checkbox beside the buttons. It picks
    /// between the `*Once` and `*Always` forms of whichever button is pressed,
    /// which is what lets two buttons carry four options.
    pub always: bool,
    reply: Reply<RequestPermissionResponse>,
}

pub struct ChatSession {
    pub id: u64,
    pub entry: settings::Agent,
    /// The directory the agent runs in, which is the project's. Kept here so
    /// a session can reconnect and file itself without asking the workspace.
    pub cwd: PathBuf,
    pub connection: Connection,
    pub items: Vec<ChatItem>,
    pub plan: Vec<(String, PlanStatus)>,
    pub permission: Option<PermissionPrompt>,
    pub commands: Vec<Command>,
    /// The modes the agent offers and the one it is in, as `session/new`
    /// reported them and every `CurrentModeUpdate` since.
    ///
    /// Belongs to the agent process rather than to the transcript, so it is
    /// filled when a connection goes live and left where it is when one goes
    /// away — the pane gates on [`Self::live`] rather than on this being empty.
    pub modes: Option<SessionModeState>,
    /// The same for config options, which is where the model lives — see
    /// `SessionConfigOptionCategory::Model`.
    pub config: Vec<SessionConfigOption>,
    /// Context spent, when the agent says.
    pub usage: Option<Usage>,
    /// What the turn in flight is counted from — see [`Self::elapsed`] and
    /// [`Self::spent`].
    flight: Option<Flight>,
    /// The agent's own name for the session, from `SessionInfoUpdate`.
    pub title: String,
    /// The name you typed, which the agent never overwrites. Two fields rather
    /// than one and a flag: whose name it is *is* the state.
    pub name: Option<String>,
    /// When the session last had something to say. Wall clock, not `Instant`,
    /// because the file has to carry it across a launch.
    pub updated: SystemTime,
    /// The agent's own id for this session — what `session/load` resumes.
    pub agent_session: Option<String>,
    /// Where the session is written, once it has anything to write.
    /// What this session is filed under, once it has been written. Minted on
    /// the first flush and not before — a session that has said nothing is not
    /// yet anything to come back to.
    pub record: Option<String>,
    /// Whether the user archived it. Typing into it clears this.
    pub closed: bool,
    pub streaming: bool,
    /// Prompts waiting for an agent to send them to: what was typed while a
    /// turn was in flight, and what a dispatched card opened the session with.
    pub queue: VecDeque<String>,
    pub transcript: transcript::State,
    _pump: Task<()>,
}

impl ChatSession {
    /// Open a session on `entry`. `seed` is its first prompt, sent as soon as
    /// the agent is up — what a dispatched card rides in on.
    pub fn connect(
        id: u64,
        entry: settings::Agent,
        cwd: PathBuf,
        seed: Option<String>,
        cx: &mut Context<Workspace>,
    ) -> Self {
        let pump = pump(id, &entry, cwd.clone(), None, cx);
        Self {
            id,
            entry,
            cwd,
            connection: Connection::Connecting,
            items: Vec::new(),
            plan: Vec::new(),
            permission: None,
            commands: Vec::new(),
            modes: None,
            config: Vec::new(),
            usage: None,
            flight: None,
            title: String::new(),
            name: None,
            updated: SystemTime::now(),
            agent_session: None,
            record: None,
            closed: false,
            streaming: false,
            queue: seed.into_iter().collect(),
            transcript: transcript::State::default(),
            _pump: pump,
        }
    }

    /// A session read back from disk. It starts idle — a launch must not spawn
    /// an agent per session — and reconnects when it is brought to the front,
    /// or when something is sent to it.
    pub fn restore(id: u64, cwd: PathBuf, entry: settings::Agent, record: Record) -> Self {
        let updated = record.at();
        Self {
            id,
            entry,
            cwd,
            connection: Connection::Idle,
            items: record.items,
            plan: Vec::new(),
            permission: None,
            commands: Vec::new(),
            modes: None,
            config: Vec::new(),
            usage: None,
            flight: None,
            title: record.title,
            name: record.name,
            updated,
            agent_session: record.session,
            record: Some(record.id),
            closed: record.closed,
            streaming: false,
            queue: VecDeque::new(),
            transcript: transcript::State::default(),
            _pump: Task::ready(()),
        }
    }

    /// When this last had something to say, as the millisecond stamp every
    /// other entry under a project carries — what the sidebar orders on.
    pub fn touched(&self) -> u128 {
        self.updated
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis()
    }

    fn to_record(&self) -> Record {
        Record {
            // The file is minted before the first write, so by the time there
            // is a record to name there is a name for it.
            id: self.record.clone().unwrap_or_default(),
            agent: self.entry.name.clone(),
            session: self.agent_session.clone(),
            title: self.title.clone(),
            name: self.name.clone(),
            updated: self
                .updated
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            closed: self.closed,
            items: self.items.clone(),
        }
    }

    /// The id this session is filed under, minted now if it has none.
    ///
    /// A dispatched card writes the link on the click that made it, and until
    /// then a session that has said nothing has no id to be written down.
    /// Bringing the mint forward costs nothing: a dispatched session carries a
    /// first prompt, so it is about to have something to file anyway.
    pub fn mint_record(&mut self) -> Option<&str> {
        if self.record.is_none() {
            self.record = fs::Project::new(&self.cwd).create_session();
        }
        self.record.as_deref()
    }

    /// Write the session out. The file is minted on the first write and not
    /// before — opening a project must not put a `.cydonia/` in it.
    pub fn flush(&mut self) {
        if self.items.is_empty() {
            return;
        }
        let store = fs::Project::new(&self.cwd);
        if self.record.is_none() {
            self.record = store.create_session();
        }
        if self.record.is_some() {
            store.save_session(&self.to_record());
        }
    }

    /// Point the session at an agent again, replaying the conversation the
    /// agent still holds when it supports `session/load`.
    pub fn resume(&mut self, cx: &mut Context<Workspace>) {
        self._pump = pump(
            self.id,
            &self.entry,
            self.cwd.clone(),
            self.agent_session.clone(),
            cx,
        );
        self.connection = Connection::Connecting;
        self.closed = false;
    }

    /// Close the connection and keep the transcript. Dropping the [`Session`]
    /// is what tears the agent process down — see its `Drop`.
    pub fn close(&mut self) {
        // Before the connection goes, not after: an agent whose stdin closes
        // half way through answering has a request it can never finish, and
        // says so on the way out.
        self.cancel();
        self.connection = Connection::Idle;
        self._pump = Task::ready(());
        self.streaming = false;
        self.closed = true;
        self.queue.clear();
        self.flush();
    }

    /// How long the turn in flight has been running.
    pub fn elapsed(&self) -> Option<Duration> {
        self.flight?.at.elapsed().ok()
    }

    /// What the turn in flight has spent, as the agent counts context.
    ///
    /// The difference rather than the total: `used` is the whole conversation,
    /// and what a running turn is costing is what it has added to it. `None`
    /// until an agent has counted at all — most do not until the first turn is
    /// answered, so a first turn shows its clock and nothing else.
    pub fn spent(&self) -> Option<u64> {
        Some(self.usage?.used.saturating_sub(self.flight?.used))
    }

    pub fn live(&self) -> bool {
        matches!(self.connection, Connection::Live(_))
    }

    /// Whether sending to it would start an agent: it is not talking to one,
    /// and one is not already on the way.
    pub fn idle(&self) -> bool {
        matches!(self.connection, Connection::Idle | Connection::Lost)
    }

    /// A session whose agent is no longer in `settings.toml` reads back but
    /// cannot reconnect — there is no command left to spawn.
    pub fn resumable(&self) -> bool {
        !self.entry.command.is_empty()
    }

    /// What to call this session: your name, else the agent's, else the
    /// agent's own name.
    pub fn label(&self) -> String {
        match (&self.name, self.title.is_empty()) {
            (Some(name), _) => name.clone(),
            (None, false) => self.title.clone(),
            (None, true) => self.entry.name.clone(),
        }
    }

    /// Send now, or queue it for whenever there is an agent to send it to —
    /// a turn in flight, or a connection still being made.
    pub fn send(&mut self, content: String) {
        if self.streaming || !self.live() {
            self.queue.push_back(content);
        } else {
            self.prompt(content);
        }
    }

    /// Send the next queued prompt, if there is one and nothing is in flight.
    pub fn drain(&mut self) {
        if self.streaming {
            return;
        }
        if let Some(next) = self.queue.pop_front() {
            self.prompt(next);
        }
    }

    fn prompt(&mut self, content: String) {
        let Connection::Live(session) = &self.connection else {
            self.queue.push_front(content);
            return;
        };
        session.prompt(&content);
        self.flight = Some(Flight {
            at: SystemTime::now(),
            used: self.usage.map_or(0, |usage| usage.used),
        });
        self.items.push(ChatItem::User(content));
        self.updated = SystemTime::now();
        self.streaming = true;
        // Before the answer, not just after it: what you said is not the
        // agent's to lose if the turn never finishes.
        self.flush();
    }

    /// Cancel the in-flight turn. A pending permission request MUST be
    /// answered `Cancelled` per spec before `session/cancel` goes out.
    pub fn cancel(&mut self) {
        if let Some(prompt) = self.permission.take() {
            prompt.reply.send(RequestPermissionResponse::cancelled());
        }
        if let Connection::Live(session) = &self.connection
            && let Err(e) = session.cancel()
        {
            self.notice(true, &format!("cancel failed: {}", acp::error_text(&e)));
        }
    }

    /// Switch the agent's mode (`session/set_mode`).
    ///
    /// Applied here as well as sent. The request is fire-and-forget and the
    /// agent's `CurrentModeUpdate` is the confirmation, so a picker that waited
    /// for it would read as a press that did nothing — and the update, when it
    /// lands, is what settles any disagreement.
    pub fn set_mode(&mut self, mode_id: &str) {
        let Connection::Live(session) = &self.connection else {
            return;
        };
        session.set_mode(mode_id);
        if let Some(modes) = &mut self.modes {
            modes.current_mode_id = mode_id.into();
        }
    }

    /// The same for a config option, which is where the model lives
    /// (`session/set_config_option`).
    pub fn set_config(&mut self, config_id: &str, value: SessionConfigOptionValue) {
        let Connection::Live(session) = &self.connection else {
            return;
        };
        session.set_config_option(config_id, value.clone());
        let found = self
            .config
            .iter_mut()
            .find(|option| &*option.id == config_id);
        let Some(option) = found else {
            return;
        };
        match (&mut option.kind, value) {
            (SessionConfigKind::Select(select), SessionConfigOptionValue::ValueId { value }) => {
                select.current_value = value;
            }
            (SessionConfigKind::Boolean(flag), SessionConfigOptionValue::Boolean { value }) => {
                flag.current_value = value;
            }
            // The agent offers one shape and was asked for the other. Its own
            // update is the answer; nothing is guessed at here.
            _ => {}
        }
    }

    /// Answer the pending permission prompt with the chosen option id.
    pub fn respond_permission(&mut self, option_id: String) {
        if let Some(prompt) = self.permission.take() {
            prompt
                .reply
                .send(RequestPermissionResponse::selected(option_id));
        }
    }

    /// Flip whether the pending prompt is answered once or for good.
    pub fn toggle_permission_always(&mut self) {
        if let Some(prompt) = &mut self.permission {
            prompt.always = !prompt.always;
        }
    }

    /// Drop the `ix`th waiting prompt — a steer taken back before the turn in
    /// flight got to it.
    pub fn unqueue(&mut self, ix: usize) {
        self.queue.remove(ix);
    }

    fn apply(&mut self, event: Event) {
        self.updated = SystemTime::now();
        match event {
            Event::Update(update) => self.apply_update(update),
            Event::Permission(request, reply) => self.open_permission(request, reply),
            Event::TurnDone(result) => {
                self.finish_thinking();
                self.streaming = false;
                match result {
                    Ok(StopReason::EndTurn) => {}
                    Ok(StopReason::Cancelled) => {
                        self.fail_running_tools();
                        self.notice(false, "cancelled");
                    }
                    Ok(StopReason::Refusal) => self.notice(false, "the agent refused to continue"),
                    Ok(StopReason::MaxTokens) => self.notice(false, "stopped: max tokens"),
                    Ok(StopReason::MaxTurnRequests) => {
                        self.notice(false, "stopped: max turn requests")
                    }
                    Ok(other) => self.notice(false, &format!("stopped: {other:?}")),
                    Err(e) => {
                        self.fail_running_tools();
                        self.notice(true, &format!("turn failed: {}", acp::error_text(&e)));
                    }
                }
                self.flush();
                self.drain();
            }
            Event::Closed => {
                self.connection = Connection::Lost;
                self.streaming = false;
                self.fail_running_tools();
                self.notice(true, "agent connection lost");
                self.flush();
            }
        }
    }

    fn apply_update(&mut self, update: SessionUpdate) {
        match update {
            SessionUpdate::AgentMessageChunk(chunk) => {
                self.finish_thinking();
                let text = content_text(&chunk.content);
                if let Some(ChatItem::Agent(body)) = self.items.last_mut() {
                    body.push_str(&text);
                } else {
                    self.items.push(ChatItem::Agent(text));
                }
            }
            SessionUpdate::AgentThoughtChunk(chunk) => {
                let text = content_text(&chunk.content);
                if let Some(ChatItem::Thinking {
                    text: body,
                    done: false,
                }) = self.items.last_mut()
                {
                    body.push_str(&text);
                } else {
                    self.items.push(ChatItem::Thinking { text, done: false });
                }
            }
            SessionUpdate::ToolCall(call) => {
                self.finish_thinking();
                self.items.push(ChatItem::Tool {
                    id: call.tool_call_id.to_string(),
                    kind: call.kind,
                    label: call.title,
                    status: tool_status(call.status),
                    output: tool_content_text(&call.content),
                });
            }
            SessionUpdate::ToolCallUpdate(update) => {
                let id = update.tool_call_id.to_string();
                let Some(ix) = self.items.iter().rposition(
                    |item| matches!(item, ChatItem::Tool { id: tool, .. } if *tool == id),
                ) else {
                    return;
                };
                let ChatItem::Tool {
                    kind,
                    label,
                    status,
                    output,
                    ..
                } = &mut self.items[ix]
                else {
                    unreachable!("rposition matched a Tool item");
                };
                if let Some(title) = update.fields.title {
                    *label = title;
                }
                if let Some(new_kind) = update.fields.kind {
                    *kind = new_kind;
                }
                if let Some(content) = update.fields.content {
                    let text = tool_content_text(&content);
                    if !text.is_empty() {
                        if !output.is_empty() {
                            output.push('\n');
                        }
                        output.push_str(&text);
                    }
                }
                if let Some(new_status) = update.fields.status {
                    *status = tool_status(new_status);
                }
            }
            SessionUpdate::SessionInfoUpdate(info) => match info.title {
                MaybeUndefined::Value(title) => self.title = title,
                MaybeUndefined::Null => self.title.clear(),
                MaybeUndefined::Undefined => {}
            },
            SessionUpdate::AvailableCommandsUpdate(cmds) => {
                self.commands = cmds
                    .available_commands
                    .into_iter()
                    .map(|c| Command {
                        name: c.name,
                        description: c.description,
                    })
                    .collect();
            }
            SessionUpdate::Plan(plan) => {
                // Plans arrive as full snapshots: replace, don't append.
                self.plan = plan
                    .entries
                    .into_iter()
                    .map(|entry| {
                        let status = match entry.status {
                            PlanEntryStatus::Completed => PlanStatus::Done,
                            PlanEntryStatus::InProgress => PlanStatus::Active,
                            _ => PlanStatus::Pending,
                        };
                        (entry.content, status)
                    })
                    .collect();
            }
            SessionUpdate::CurrentModeUpdate(update) => {
                if let Some(modes) = &mut self.modes {
                    modes.current_mode_id = update.current_mode_id;
                }
            }
            // Reported whole, like the plan: replace, don't merge.
            SessionUpdate::ConfigOptionUpdate(update) => self.config = update.config_options,
            SessionUpdate::UsageUpdate(update) => {
                self.usage = Some(Usage {
                    used: update.used,
                    size: update.size,
                });
            }
            // We echo the user's message locally.
            SessionUpdate::UserMessageChunk(_) => {}
            _ => {}
        }
    }

    fn open_permission(
        &mut self,
        request: RequestPermissionRequest,
        reply: Reply<RequestPermissionResponse>,
    ) {
        let options: Vec<Choice> = request
            .options
            .into_iter()
            .map(|opt| Choice {
                id: opt.option_id.to_string(),
                name: opt.name,
                kind: opt.kind,
            })
            .collect();
        if options.is_empty() {
            reply.send(RequestPermissionResponse::cancelled());
            return;
        }
        // A replaced prompt must still be answered — an unanswered
        // reply hangs the agent.
        if let Some(previous) = self.permission.take() {
            previous.reply.send(RequestPermissionResponse::cancelled());
        }
        let title = request
            .tool_call
            .fields
            .title
            .clone()
            .unwrap_or_else(|| "Permission required".to_owned());
        self.permission = Some(PermissionPrompt {
            title,
            options,
            always: false,
            reply,
        });
    }

    pub(crate) fn notice(&mut self, failed: bool, text: &str) {
        self.items.push(ChatItem::Notice {
            text: text.to_owned(),
            failed,
        });
    }

    fn finish_thinking(&mut self) {
        if let Some(ChatItem::Thinking { done, .. }) = self.items.last_mut() {
            *done = true;
        }
    }

    fn fail_running_tools(&mut self) {
        for item in &mut self.items {
            if let ChatItem::Tool { status, .. } = item
                && *status == ToolStatus::Running
            {
                *status = ToolStatus::Failure;
            }
        }
    }
}

fn tool_status(status: ToolCallStatus) -> ToolStatus {
    match status {
        ToolCallStatus::Completed => ToolStatus::Success,
        ToolCallStatus::Failed => ToolStatus::Failure,
        _ => ToolStatus::Running,
    }
}

fn tool_content_text(content: &[ToolCallContent]) -> String {
    content
        .iter()
        .map(|c| match c {
            ToolCallContent::Content { content } => content_text(content),
            ToolCallContent::Diff(diff) => format!("edited {}", diff.path.display()),
            ToolCallContent::Terminal { .. } => "[terminal]".to_owned(),
            _ => "[content]".to_owned(),
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn content_text(block: &ContentBlock) -> String {
    match block {
        ContentBlock::Text(text) => text.text.clone(),
        _ => "[non-text content]".to_owned(),
    }
}

/// Drain the ACP event channel into the session, for as long as there is one.
///
/// `previous` is the agent's own session id: with it set and the agent
/// capable, the conversation is loaded rather than started over, so the turn
/// that follows carries everything said before it.
fn pump(
    id: u64,
    entry: &settings::Agent,
    cwd: PathBuf,
    previous: Option<String>,
    cx: &mut Context<Workspace>,
) -> Task<()> {
    let entry = entry.clone();
    let conn = acp::runtime().spawn(async move {
        Session::spawn(
            &entry,
            Launch {
                previous,
                ..Launch::new(cwd)
            },
        )
        .await
    });

    cx.spawn(async move |this, cx| {
        let opened = conn
            .await
            .unwrap_or_else(|e| Err(anyhow!("the connection task panicked: {e}")));
        let (session, mut events) = match opened {
            Ok(pair) => pair,
            Err(e) => {
                let _ = this.update(cx, |workspace, cx| {
                    workspace.with_session(id, cx, |chat| {
                        chat.connection = Connection::Lost;
                        chat.notice(true, &format!("connection failed: {e:#}"));
                    });
                });
                return;
            }
        };

        if this
            .update(cx, |workspace, cx| {
                workspace.with_session(id, cx, |chat| {
                    chat.agent_session = Some(session.session_id.to_string());
                    // What this agent lets you switch, read off `session/new`
                    // before the session is boxed away. Every change after
                    // this arrives as an update.
                    chat.modes = session.response.modes.clone();
                    chat.config = session.response.config_options.clone().unwrap_or_default();
                    chat.connection = Connection::Live(Box::new(session));
                });
                workspace.session_connected(id, cx);
            })
            .is_err()
        {
            return;
        }

        while let Some(event) = events.recv().await {
            let mut batch = vec![event];
            while let Ok(event) = events.try_recv() {
                batch.push(event);
            }
            let streaming = this.update(cx, |workspace, cx| {
                workspace.with_session(id, cx, |chat| {
                    for event in batch {
                        chat.apply(event);
                    }
                });
                workspace.session(id).is_some_and(|chat| chat.streaming)
            });
            match streaming {
                Ok(true) => cx.background_executor().timer(STREAM_FRAME).await,
                Ok(false) => {}
                Err(_) => return,
            }
        }
    })
}