Allow displaying diffs in the ui. (#8)
Reviewed-on: #8 Co-authored-by: Drew Galbraith <drew@tiramisu.one> Co-committed-by: Drew Galbraith <drew@tiramisu.one>
This commit is contained in:
parent
af080710cc
commit
5a49cba1e6
9 changed files with 481 additions and 114 deletions
|
|
@ -3,20 +3,25 @@
|
|||
use tokio::sync::mpsc;
|
||||
use tracing::debug;
|
||||
|
||||
use super::AppState;
|
||||
use super::{AppState, DisplayMessage};
|
||||
use crate::core::types::{Role, StampedEvent, UIEvent};
|
||||
use crate::tui::tool_display;
|
||||
|
||||
/// Drain all pending [`UIEvent`]s from `event_rx` and apply them to `state`.
|
||||
///
|
||||
/// This is non-blocking: it processes all currently-available events and returns
|
||||
/// immediately when the channel is empty.
|
||||
///
|
||||
/// Tool events use in-place replacement: when a `ToolExecuting` or `ToolResult`
|
||||
/// arrives, the handler searches `state.messages` for an existing entry with the
|
||||
/// same `tool_use_id` and replaces its content rather than appending a new row.
|
||||
///
|
||||
/// | Event | Effect |
|
||||
/// |------------------------|------------------------------------------------------------|
|
||||
/// | `StreamDelta(s)` | Append `s` to last message if it's `Assistant`; else push |
|
||||
/// | `ToolApprovalRequest` | Set `pending_approval` in state |
|
||||
/// | `ToolExecuting` | Display tool execution info |
|
||||
/// | `ToolResult` | Display tool result |
|
||||
/// | `ToolApprovalRequest` | Push inline message with approval prompt, set pending |
|
||||
/// | `ToolExecuting` | Replace approval message in-place (or push new) |
|
||||
/// | `ToolResult` | Replace executing message in-place (or push new) |
|
||||
/// | `TurnComplete` | No structural change; logged at debug level |
|
||||
/// | `Error(msg)` | Push `(Assistant, "[error] {msg}")` |
|
||||
pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<StampedEvent>, state: &mut AppState) {
|
||||
|
|
@ -32,43 +37,57 @@ pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<StampedEvent>, state
|
|||
}
|
||||
match stamped.event {
|
||||
UIEvent::StreamDelta(chunk) => {
|
||||
if let Some((Role::Assistant, content)) = state.messages.last_mut() {
|
||||
content.push_str(&chunk);
|
||||
if let Some(msg) = state.messages.last_mut() {
|
||||
if msg.role == Role::Assistant && msg.tool_use_id.is_none() {
|
||||
msg.content.push_str(&chunk);
|
||||
} else {
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content: chunk,
|
||||
tool_use_id: None,
|
||||
});
|
||||
}
|
||||
} else {
|
||||
state.messages.push((Role::Assistant, chunk));
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content: chunk,
|
||||
tool_use_id: None,
|
||||
});
|
||||
}
|
||||
state.content_changed = true;
|
||||
}
|
||||
UIEvent::ToolApprovalRequest {
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
input_summary,
|
||||
display,
|
||||
} => {
|
||||
state.pending_approval = Some(PendingApproval {
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
input_summary,
|
||||
let mut content = tool_display::format_executing(&tool_name, &display);
|
||||
content.push_str("\n[y] approve [n] deny");
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content,
|
||||
tool_use_id: Some(tool_use_id.clone()),
|
||||
});
|
||||
state.pending_approval = Some(PendingApproval { tool_use_id });
|
||||
state.content_changed = true;
|
||||
}
|
||||
UIEvent::ToolExecuting {
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
input_summary,
|
||||
display,
|
||||
} => {
|
||||
state
|
||||
.messages
|
||||
.push((Role::Assistant, format!("[{tool_name}] {input_summary}")));
|
||||
let content = tool_display::format_executing(&tool_name, &display);
|
||||
replace_or_push(state, &tool_use_id, content);
|
||||
state.content_changed = true;
|
||||
}
|
||||
UIEvent::ToolResult {
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
output_summary,
|
||||
display,
|
||||
is_error,
|
||||
} => {
|
||||
let prefix = if is_error { "error" } else { "result" };
|
||||
state.messages.push((
|
||||
Role::Assistant,
|
||||
format!("[{tool_name} {prefix}] {output_summary}"),
|
||||
));
|
||||
let content = tool_display::format_result(&tool_name, &display, is_error);
|
||||
replace_or_push(state, &tool_use_id, content);
|
||||
state.content_changed = true;
|
||||
}
|
||||
UIEvent::TurnComplete => {
|
||||
|
|
@ -78,26 +97,46 @@ pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<StampedEvent>, state
|
|||
state.network_allowed = allowed;
|
||||
}
|
||||
UIEvent::Error(msg) => {
|
||||
state
|
||||
.messages
|
||||
.push((Role::Assistant, format!("[error] {msg}")));
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content: format!("[error] {msg}"),
|
||||
tool_use_id: None,
|
||||
});
|
||||
state.content_changed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the message with the given `tool_use_id` and replace its content,
|
||||
/// or push a new message if not found.
|
||||
fn replace_or_push(state: &mut AppState, tool_use_id: &str, content: String) {
|
||||
if let Some(msg) = state
|
||||
.messages
|
||||
.iter_mut()
|
||||
.rev()
|
||||
.find(|m| m.tool_use_id.as_deref() == Some(tool_use_id))
|
||||
{
|
||||
msg.content = content;
|
||||
} else {
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content,
|
||||
tool_use_id: Some(tool_use_id.to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// A pending tool approval request waiting for user input.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PendingApproval {
|
||||
pub tool_use_id: String,
|
||||
pub tool_name: String,
|
||||
pub input_summary: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::core::types::ToolDisplay;
|
||||
|
||||
/// Wrap a [`UIEvent`] in a [`StampedEvent`] at epoch 0 for tests.
|
||||
fn stamp(event: UIEvent) -> StampedEvent {
|
||||
|
|
@ -108,63 +147,86 @@ mod tests {
|
|||
async fn drain_appends_to_existing_assistant() {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
|
||||
let mut state = AppState::new();
|
||||
state.messages.push((Role::Assistant, "hello".to_string()));
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content: "hello".to_string(),
|
||||
tool_use_id: None,
|
||||
});
|
||||
tx.send(stamp(UIEvent::StreamDelta(" world".to_string())))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
drain_ui_events(&mut rx, &mut state);
|
||||
assert_eq!(state.messages.last().unwrap().1, "hello world");
|
||||
assert_eq!(state.messages.last().unwrap().content, "hello world");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_creates_assistant_on_user_last() {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
|
||||
let mut state = AppState::new();
|
||||
state.messages.push((Role::User, "hi".to_string()));
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::User,
|
||||
content: "hi".to_string(),
|
||||
tool_use_id: None,
|
||||
});
|
||||
tx.send(stamp(UIEvent::StreamDelta("hello".to_string())))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
drain_ui_events(&mut rx, &mut state);
|
||||
assert_eq!(state.messages.len(), 2);
|
||||
assert_eq!(state.messages[1].0, Role::Assistant);
|
||||
assert_eq!(state.messages[1].1, "hello");
|
||||
assert_eq!(state.messages[1].role, Role::Assistant);
|
||||
assert_eq!(state.messages[1].content, "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_tool_approval_sets_pending() {
|
||||
async fn drain_tool_approval_sets_pending_and_adds_message() {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
|
||||
let mut state = AppState::new();
|
||||
tx.send(stamp(UIEvent::ToolApprovalRequest {
|
||||
tool_use_id: "t1".to_string(),
|
||||
tool_name: "write_file".to_string(),
|
||||
input_summary: "path: foo.txt".to_string(),
|
||||
tool_name: "shell_exec".to_string(),
|
||||
display: ToolDisplay::ShellExec {
|
||||
command: "cargo test".to_string(),
|
||||
},
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
drain_ui_events(&mut rx, &mut state);
|
||||
assert!(state.pending_approval.is_some());
|
||||
let approval = state.pending_approval.unwrap();
|
||||
assert_eq!(approval.tool_name, "write_file");
|
||||
assert_eq!(state.pending_approval.as_ref().unwrap().tool_use_id, "t1");
|
||||
// Message should be inline with approval prompt.
|
||||
assert_eq!(state.messages.len(), 1);
|
||||
assert!(state.messages[0].content.contains("[y] approve"));
|
||||
assert_eq!(state.messages[0].tool_use_id.as_deref(), Some("t1"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_tool_result_adds_message() {
|
||||
async fn drain_tool_result_replaces_existing_message() {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
|
||||
let mut state = AppState::new();
|
||||
// Simulate an existing executing message.
|
||||
state.messages.push(DisplayMessage {
|
||||
role: Role::Assistant,
|
||||
content: "$ cargo test".to_string(),
|
||||
tool_use_id: Some("t1".to_string()),
|
||||
});
|
||||
tx.send(stamp(UIEvent::ToolResult {
|
||||
tool_name: "read_file".to_string(),
|
||||
output_summary: "file contents...".to_string(),
|
||||
tool_use_id: "t1".to_string(),
|
||||
tool_name: "shell_exec".to_string(),
|
||||
display: ToolDisplay::ShellExec {
|
||||
command: "cargo test\nok".to_string(),
|
||||
},
|
||||
is_error: false,
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
drain_ui_events(&mut rx, &mut state);
|
||||
// Should replace, not append.
|
||||
assert_eq!(state.messages.len(), 1);
|
||||
assert!(state.messages[0].1.contains("read_file result"));
|
||||
assert!(state.messages[0].content.contains("cargo test"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -172,14 +234,12 @@ mod tests {
|
|||
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
|
||||
let mut state = AppState::new();
|
||||
state.epoch = 2;
|
||||
// Event from epoch 1 should be discarded.
|
||||
tx.send(StampedEvent {
|
||||
epoch: 1,
|
||||
event: UIEvent::StreamDelta("ghost".to_string()),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
// Event from epoch 2 should be accepted.
|
||||
tx.send(StampedEvent {
|
||||
epoch: 2,
|
||||
event: UIEvent::StreamDelta("real".to_string()),
|
||||
|
|
@ -189,7 +249,7 @@ mod tests {
|
|||
drop(tx);
|
||||
drain_ui_events(&mut rx, &mut state);
|
||||
assert_eq!(state.messages.len(), 1);
|
||||
assert_eq!(state.messages[0].1, "real");
|
||||
assert_eq!(state.messages[0].content, "real");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue