notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

shared.rs (7256B)


      1 //! Shared utilities used by multiple AI backend implementations.
      2 
      3 use crate::auto_accept::AutoAcceptRules;
      4 use crate::backend::tool_summary::{format_tool_summary, truncate_output};
      5 use crate::file_update::FileUpdate;
      6 use crate::messages::{DaveApiResponse, ExecutedTool, PendingPermission, PermissionRequest};
      7 use crate::Message;
      8 use claude_agent_sdk_rs::PermissionMode;
      9 use std::sync::mpsc;
     10 use tokio::sync::mpsc as tokio_mpsc;
     11 use tokio::sync::oneshot;
     12 use uuid::Uuid;
     13 
     14 /// Commands sent to a session's actor task.
     15 ///
     16 /// Used identically by the Claude and Codex backends.
     17 pub(crate) enum SessionCommand {
     18     Query {
     19         prompt: String,
     20         response_tx: mpsc::Sender<DaveApiResponse>,
     21         ctx: egui::Context,
     22     },
     23     /// Interrupt the current query - stops the stream but preserves session
     24     Interrupt {
     25         ctx: egui::Context,
     26     },
     27     /// Set the permission mode (Default or Plan)
     28     SetPermissionMode {
     29         mode: PermissionMode,
     30         ctx: egui::Context,
     31     },
     32     /// Trigger manual context compaction
     33     Compact {
     34         response_tx: mpsc::Sender<DaveApiResponse>,
     35         ctx: egui::Context,
     36     },
     37     Shutdown,
     38 }
     39 
     40 /// Handle kept by a backend to communicate with its session actor.
     41 pub(crate) struct SessionHandle {
     42     pub command_tx: tokio_mpsc::Sender<SessionCommand>,
     43 }
     44 
     45 /// Convert our messages to a prompt string for the AI backend.
     46 ///
     47 /// Includes the system message (if any) followed by the conversation
     48 /// history formatted as `Human:` / `Assistant:` turns. Tool-related,
     49 /// error, permission, compaction and subagent messages are skipped.
     50 pub fn messages_to_prompt(messages: &[Message]) -> String {
     51     let mut prompt = String::new();
     52 
     53     // Include system message if present
     54     for msg in messages {
     55         if let Message::System(content) = msg {
     56             prompt.push_str(content);
     57             prompt.push_str("\n\n");
     58             break;
     59         }
     60     }
     61 
     62     // Format conversation history
     63     for msg in messages {
     64         match msg {
     65             Message::System(_) => {} // Already handled
     66             Message::User(content) => {
     67                 prompt.push_str("Human: ");
     68                 prompt.push_str(content);
     69                 prompt.push_str("\n\n");
     70             }
     71             Message::Assistant(content) => {
     72                 prompt.push_str("Assistant: ");
     73                 prompt.push_str(content.text());
     74                 prompt.push_str("\n\n");
     75             }
     76             Message::ToolCalls(_)
     77             | Message::ToolResponse(_)
     78             | Message::Error(_)
     79             | Message::PermissionRequest(_)
     80             | Message::CompactionComplete(_)
     81             | Message::Subagent(_) => {}
     82         }
     83     }
     84 
     85     prompt
     86 }
     87 
     88 /// Collect all trailing user messages and join them.
     89 ///
     90 /// When multiple messages are queued, they're all sent as one prompt
     91 /// so the AI sees everything at once instead of one at a time.
     92 pub fn get_pending_user_messages(messages: &[Message]) -> String {
     93     let mut trailing: Vec<&str> = messages
     94         .iter()
     95         .rev()
     96         .take_while(|m| matches!(m, Message::User(_)))
     97         .filter_map(|m| match m {
     98             Message::User(content) => Some(content.as_str()),
     99             _ => None,
    100         })
    101         .collect();
    102     trailing.reverse();
    103     trailing.join("\n")
    104 }
    105 
    106 /// Remove a completed subagent from the stack and notify the UI.
    107 pub fn complete_subagent(
    108     task_id: &str,
    109     result_text: &str,
    110     subagent_stack: &mut Vec<String>,
    111     response_tx: &mpsc::Sender<DaveApiResponse>,
    112     ctx: &egui::Context,
    113 ) {
    114     subagent_stack.retain(|id| id != task_id);
    115     let _ = response_tx.send(DaveApiResponse::SubagentCompleted {
    116         task_id: task_id.to_string(),
    117         result: truncate_output(result_text, 2000),
    118     });
    119     ctx.request_repaint();
    120 }
    121 
    122 /// Build an [`ExecutedTool`] from a completed tool call and send it
    123 /// to the UI along with a repaint request.
    124 pub fn send_tool_result(
    125     tool_name: &str,
    126     tool_input: &serde_json::Value,
    127     result_value: &serde_json::Value,
    128     file_update: Option<FileUpdate>,
    129     subagent_stack: &[String],
    130     response_tx: &mpsc::Sender<DaveApiResponse>,
    131     ctx: &egui::Context,
    132 ) {
    133     let summary = format_tool_summary(tool_name, tool_input, result_value);
    134     let parent_task_id = subagent_stack.last().cloned();
    135     let tool_result = ExecutedTool {
    136         tool_name: tool_name.to_string(),
    137         summary,
    138         parent_task_id,
    139         file_update,
    140     };
    141     let _ = response_tx.send(DaveApiResponse::ToolResult(tool_result));
    142     ctx.request_repaint();
    143 }
    144 
    145 /// Check auto-accept rules for a tool invocation.
    146 ///
    147 /// Returns `true` (and logs) when the tool should be silently
    148 /// accepted without asking the user.
    149 pub fn should_auto_accept(tool_name: &str, tool_input: &serde_json::Value) -> bool {
    150     let rules = AutoAcceptRules::default();
    151     let accepted = rules.should_auto_accept(tool_name, tool_input);
    152     if accepted {
    153         tracing::debug!("Auto-accepting {}: matched auto-accept rule", tool_name);
    154     }
    155     accepted
    156 }
    157 
    158 /// Build a [`PermissionRequest`] + [`PendingPermission`], send it to
    159 /// the UI via `response_tx`, and return the oneshot receiver the
    160 /// caller can `await` to get the user's decision.
    161 ///
    162 /// Returns `None` if the UI channel is closed (the request could not
    163 /// be delivered).
    164 pub fn forward_permission_to_ui(
    165     tool_name: &str,
    166     tool_input: serde_json::Value,
    167     response_tx: &mpsc::Sender<DaveApiResponse>,
    168     ctx: &egui::Context,
    169 ) -> Option<oneshot::Receiver<crate::messages::PermissionResponse>> {
    170     let request_id = Uuid::new_v4();
    171     let (ui_resp_tx, ui_resp_rx) = oneshot::channel();
    172 
    173     let cached_plan = if tool_name == "ExitPlanMode" {
    174         tool_input
    175             .get("plan")
    176             .and_then(|v| v.as_str())
    177             .map(crate::messages::ParsedMarkdown::parse)
    178     } else {
    179         None
    180     };
    181 
    182     let request = PermissionRequest {
    183         id: request_id,
    184         tool_name: tool_name.to_string(),
    185         tool_input,
    186         response: None,
    187         answer_summary: None,
    188         cached_plan,
    189     };
    190 
    191     let pending = PendingPermission {
    192         request,
    193         response_tx: ui_resp_tx,
    194     };
    195 
    196     if response_tx
    197         .send(DaveApiResponse::PermissionRequest(pending))
    198         .is_err()
    199     {
    200         tracing::error!("Failed to send permission request to UI");
    201         return None;
    202     }
    203 
    204     ctx.request_repaint();
    205     Some(ui_resp_rx)
    206 }
    207 
    208 /// Decide which prompt to send based on whether we're resuming a
    209 /// session and how many user messages exist.
    210 ///
    211 /// - Resumed sessions always send just the pending messages (the
    212 ///   backend already has the full conversation context).
    213 /// - New sessions send the full prompt on the first message, then
    214 ///   only pending messages for subsequent turns.
    215 pub fn prepare_prompt(messages: &[Message], resume_session_id: &Option<String>) -> String {
    216     if resume_session_id.is_some() {
    217         get_pending_user_messages(messages)
    218     } else {
    219         let is_first_message = messages
    220             .iter()
    221             .filter(|m| matches!(m, Message::User(_)))
    222             .count()
    223             == 1;
    224         if is_first_message {
    225             messages_to_prompt(messages)
    226         } else {
    227             get_pending_user_messages(messages)
    228         }
    229     }
    230 }