zed/crates/agent_ui/src/terminal_codegen.rs

use crate::inline_prompt_editor::CodegenStatus;
use futures::{SinkExt, StreamExt, channel::mpsc};
use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task};
use language_model::{ConfiguredModel, LanguageModelRegistry, LanguageModelRequest};
use language_models::provider::anthropic::telemetry::{
    AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,
};
use std::time::Instant;
use terminal::Terminal;
use uuid::Uuid;

pub struct TerminalCodegen {
    pub status: CodegenStatus,
    terminal: Entity<Terminal>,
    generation: Task<()>,
    pub message_id: Option<String>,
    transaction: Option<TerminalTransaction>,
    session_id: Uuid,
}

impl EventEmitter<CodegenEvent> for TerminalCodegen {}

impl TerminalCodegen {
    pub fn new(terminal: Entity<Terminal>, session_id: Uuid) -> Self {
        Self {
            terminal,
            status: CodegenStatus::Idle,
            generation: Task::ready(()),
            message_id: None,
            transaction: None,
            session_id,
        }
    }

    pub fn session_id(&self) -> Uuid {
        self.session_id
    }

    pub fn start(&mut self, prompt_task: Task<LanguageModelRequest>, cx: &mut Context<Self>) {
        let Some(ConfiguredModel { model, .. }) =
            LanguageModelRegistry::read_global(cx).inline_assistant_model()
        else {
            return;
        };

        let anthropic_reporter = AnthropicEventReporter::new(&model, cx);
        let session_id = self.session_id;
        let model_telemetry_id = model.telemetry_id();
        let model_provider_id = model.provider_id().to_string();

        self.status = CodegenStatus::Pending;
        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
        self.generation = cx.spawn(async move |this, cx| {
            let prompt = prompt_task.await;
            let response = model.stream_completion_text(prompt, cx).await;
            let generate = async {
                let message_id = response
                    .as_ref()
                    .ok()
                    .and_then(|response| response.message_id.clone());

                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);

                let task = cx.background_spawn({
                    let message_id = message_id.clone();
                    let anthropic_reporter = anthropic_reporter.clone();
                    async move {
                        let mut response_latency = None;
                        let request_start = Instant::now();
                        let task = async {
                            let mut chunks = response?.stream;
                            while let Some(chunk) = chunks.next().await {
                                if response_latency.is_none() {
                                    response_latency = Some(request_start.elapsed());
                                }
                                let chunk = chunk?;
                                hunks_tx.send(chunk).await?;
                            }

                            anyhow::Ok(())
                        };

                        let result = task.await;

                        let error_message = result.as_ref().err().map(|error| error.to_string());

                        telemetry::event!(
                            "Assistant Responded",
                            session_id = session_id.to_string(),
                            kind = "inline_terminal",
                            phase = "response",
                            model = model_telemetry_id,
                            model_provider = model_provider_id,
                            language_name = Option::<&str>::None,
                            message_id = message_id,
                            response_latency = response_latency,
                            error_message = error_message,
                        );

                        anthropic_reporter.report(AnthropicEventData {
                            completion_type: AnthropicCompletionType::Terminal,
                            event: AnthropicEventType::Response,
                            language_name: None,
                            message_id,
                        });

                        result?;
                        anyhow::Ok(())
                    }
                });

                this.update(cx, |this, _| {
                    this.message_id = message_id;
                })?;

                while let Some(hunk) = hunks_rx.next().await {
                    this.update(cx, |this, cx| {
                        if let Some(transaction) = &mut this.transaction {
                            transaction.push(hunk, cx);
                            cx.notify();
                        }
                    })?;
                }

                task.await?;
                anyhow::Ok(())
            };

            let result = generate.await;

            this.update(cx, |this, cx| {
                if let Err(error) = result {
                    this.status = CodegenStatus::Error(error);
                } else {
                    this.status = CodegenStatus::Done;
                }
                cx.emit(CodegenEvent::Finished);
                cx.notify();
            })
            .ok();
        });
        cx.notify();
    }

    pub fn completion(&self) -> Option<String> {
        self.transaction
            .as_ref()
            .map(|transaction| transaction.completion.clone())
    }

    pub fn stop(&mut self, cx: &mut Context<Self>) {
        self.status = CodegenStatus::Done;
        self.generation = Task::ready(());
        cx.emit(CodegenEvent::Finished);
        cx.notify();
    }

    pub fn complete(&mut self, cx: &mut Context<Self>) {
        if let Some(transaction) = self.transaction.take() {
            transaction.complete(cx);
        }
    }

    pub fn undo(&mut self, cx: &mut Context<Self>) {
        if let Some(transaction) = self.transaction.take() {
            transaction.undo(cx);
        }
    }
}

#[derive(Copy, Clone, Debug)]
pub enum CodegenEvent {
    Finished,
}

#[cfg(not(target_os = "windows"))]
pub const CLEAR_INPUT: &str = "\x15";
#[cfg(target_os = "windows")]
pub const CLEAR_INPUT: &str = "\x03";
const CARRIAGE_RETURN: &str = "\x0d";

struct TerminalTransaction {
    completion: String,
    terminal: Entity<Terminal>,
}

impl TerminalTransaction {
    pub fn start(terminal: Entity<Terminal>) -> Self {
        Self {
            completion: String::new(),
            terminal,
        }
    }

    pub fn push(&mut self, hunk: String, cx: &mut App) {
        // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
        let input = Self::sanitize_input(hunk);
        self.completion.push_str(&input);
        self.terminal
            .update(cx, |terminal, _| terminal.input(input.into_bytes()));
    }

    pub fn undo(self, cx: &mut App) {
        self.terminal
            .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.as_bytes()));
    }

    pub fn complete(self, cx: &mut App) {
        self.terminal
            .update(cx, |terminal, _| terminal.input(CARRIAGE_RETURN.as_bytes()));
    }

    fn sanitize_input(mut input: String) -> String {
        input.retain(|c| c != '\r' && c != '\n');
        input
    }
}