refact
This commit is contained in:
235
src/ai.rs
235
src/ai.rs
@@ -1,33 +1,20 @@
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::sync::mpsc;
|
||||
use std::thread;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionStatus {
|
||||
pub id: usize,
|
||||
pub state: String,
|
||||
pub label: String,
|
||||
}
|
||||
|
||||
/// Messages sent from the reader thread to the main display loop.
|
||||
pub enum OutputEvent {
|
||||
/// Start of a new AI response.
|
||||
StreamStart,
|
||||
/// A chunk of text to display immediately.
|
||||
StreamChunk(String),
|
||||
/// End of the current response.
|
||||
StreamEnd,
|
||||
}
|
||||
|
||||
pub struct ClaudeManager {
|
||||
stdin: Arc<Mutex<std::process::ChildStdin>>,
|
||||
status: Arc<Mutex<Vec<SessionStatus>>>,
|
||||
stdin: std::process::ChildStdin,
|
||||
state: Arc<Mutex<String>>,
|
||||
output_rx: mpsc::Receiver<OutputEvent>,
|
||||
id_tx: mpsc::Sender<usize>,
|
||||
next_id: usize,
|
||||
child_pid: u32,
|
||||
}
|
||||
|
||||
@@ -45,119 +32,81 @@ impl ClaudeManager {
|
||||
.map_err(|e| format!("failed to start claude: {}", e))?;
|
||||
|
||||
let child_pid = child.id();
|
||||
let stdin = child.stdin.take()
|
||||
.ok_or("failed to capture claude stdin")?;
|
||||
let stdin = Arc::new(Mutex::new(stdin));
|
||||
let child_stdout = child.stdout.take()
|
||||
.ok_or("failed to capture claude stdout")?;
|
||||
let stdin = child.stdin.take().ok_or("failed to capture stdin")?;
|
||||
let stdout = child.stdout.take().ok_or("failed to capture stdout")?;
|
||||
|
||||
let status: Arc<Mutex<Vec<SessionStatus>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let (output_tx, output_rx) = mpsc::channel();
|
||||
let (id_tx, id_rx) = mpsc::channel::<usize>();
|
||||
let state = Arc::new(Mutex::new("idle".to_string()));
|
||||
let (tx, rx) = mpsc::channel();
|
||||
let state_bg = state.clone();
|
||||
|
||||
let status_clone = status.clone();
|
||||
|
||||
// Background thread: read Claude's stdout, stream responses
|
||||
thread::spawn(move || {
|
||||
Self::reader_loop(child, child_stdout, id_rx, output_tx, status_clone);
|
||||
Self::reader_loop(child, stdout, tx, state_bg);
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
stdin,
|
||||
status,
|
||||
output_rx,
|
||||
id_tx,
|
||||
next_id: 1,
|
||||
child_pid,
|
||||
})
|
||||
Ok(Self { stdin, state, output_rx: rx, child_pid })
|
||||
}
|
||||
|
||||
fn reader_loop(
|
||||
mut child: Child,
|
||||
mut child: std::process::Child,
|
||||
stdout: std::process::ChildStdout,
|
||||
id_rx: mpsc::Receiver<usize>,
|
||||
output_tx: mpsc::Sender<OutputEvent>,
|
||||
status: Arc<Mutex<Vec<SessionStatus>>>,
|
||||
tx: mpsc::Sender<OutputEvent>,
|
||||
state: Arc<Mutex<String>>,
|
||||
) {
|
||||
let reader = BufReader::new(stdout);
|
||||
let mut current_id: Option<usize> = None;
|
||||
let mut stream_started = false;
|
||||
let mut last_text_len: usize = 0;
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = match line {
|
||||
Ok(l) => l,
|
||||
Err(_) => break,
|
||||
};
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
for line in reader.lines().flatten() {
|
||||
if line.trim().is_empty() { continue; }
|
||||
let json: Value = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let msg_type = json.get("type").and_then(|t| t.as_str());
|
||||
match msg_type {
|
||||
Some("system") => {}
|
||||
match json.get("type").and_then(|t| t.as_str()) {
|
||||
Some("assistant") => {
|
||||
// Pick up the session ID for this response
|
||||
if current_id.is_none() {
|
||||
current_id = id_rx.try_recv().ok();
|
||||
}
|
||||
|
||||
// Stream text chunks
|
||||
if let Some(content) = json.pointer("/message/content").and_then(|c| c.as_array()) {
|
||||
for item in content {
|
||||
match item.get("type").and_then(|t| t.as_str()) {
|
||||
Some("text") => {
|
||||
if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
|
||||
// Send stream start on first chunk
|
||||
if !stream_started {
|
||||
let _ = output_tx.send(OutputEvent::StreamStart);
|
||||
stream_started = true;
|
||||
last_text_len = 0;
|
||||
}
|
||||
// Only send the new part (delta)
|
||||
if text.len() > last_text_len {
|
||||
let delta = &text[last_text_len..];
|
||||
let _ = output_tx.send(OutputEvent::StreamChunk(delta.to_string()));
|
||||
last_text_len = text.len();
|
||||
}
|
||||
let content = json.pointer("/message/content").and_then(|c| c.as_array());
|
||||
for item in content.into_iter().flatten() {
|
||||
match item.get("type").and_then(|t| t.as_str()) {
|
||||
Some("text") => {
|
||||
if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
|
||||
if !stream_started {
|
||||
let _ = tx.send(OutputEvent::StreamStart);
|
||||
stream_started = true;
|
||||
last_text_len = 0;
|
||||
}
|
||||
if text.len() > last_text_len {
|
||||
let _ = tx.send(OutputEvent::StreamChunk(
|
||||
text[last_text_len..].to_string(),
|
||||
));
|
||||
last_text_len = text.len();
|
||||
}
|
||||
Self::update_status(&status, current_id, "responding...");
|
||||
}
|
||||
Some("tool_use") => {
|
||||
let name = item.get("name")
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("tool");
|
||||
Self::update_status(&status, current_id, &format!("running: {}...", name));
|
||||
}
|
||||
_ => {}
|
||||
Self::set_state(&state, "responding...");
|
||||
}
|
||||
Some("tool_use") => {
|
||||
let name = item.get("name")
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("tool");
|
||||
Self::set_state(&state, &format!("running: {name}..."));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("result") => {
|
||||
// If we never streamed text, send the result field
|
||||
if !stream_started {
|
||||
if let Some(result) = json.get("result").and_then(|r| r.as_str()) {
|
||||
if !result.is_empty() {
|
||||
let _ = output_tx.send(OutputEvent::StreamStart);
|
||||
let _ = output_tx.send(OutputEvent::StreamChunk(result.to_string()));
|
||||
if let Some(text) = json.get("result").and_then(|r| r.as_str()) {
|
||||
if !text.is_empty() {
|
||||
let _ = tx.send(OutputEvent::StreamStart);
|
||||
let _ = tx.send(OutputEvent::StreamChunk(text.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = output_tx.send(OutputEvent::StreamEnd);
|
||||
let _ = tx.send(OutputEvent::StreamEnd);
|
||||
stream_started = false;
|
||||
last_text_len = 0;
|
||||
|
||||
// Remove session from status
|
||||
if let Some(id) = current_id.take() {
|
||||
if let Ok(mut st) = status.lock() {
|
||||
st.retain(|s| s.id != id);
|
||||
}
|
||||
}
|
||||
Self::set_state(&state, "idle");
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
@@ -165,107 +114,35 @@ impl ClaudeManager {
|
||||
let _ = child.wait();
|
||||
}
|
||||
|
||||
fn update_status(
|
||||
status: &Arc<Mutex<Vec<SessionStatus>>>,
|
||||
id: Option<usize>,
|
||||
state: &str,
|
||||
) {
|
||||
let id = match id {
|
||||
Some(id) => id,
|
||||
None => return,
|
||||
};
|
||||
if let Ok(mut st) = status.lock() {
|
||||
if let Some(s) = st.iter_mut().find(|s| s.id == id) {
|
||||
s.state = state.to_string();
|
||||
}
|
||||
}
|
||||
fn set_state(state: &Arc<Mutex<String>>, s: &str) {
|
||||
if let Ok(mut st) = state.lock() { *st = s.to_string(); }
|
||||
}
|
||||
|
||||
/// Send a message to Claude (non-blocking). Returns session ID.
|
||||
pub fn send(&mut self, input: &str) -> usize {
|
||||
// Find lowest available ID (reuse completed IDs)
|
||||
let id = if let Ok(st) = self.status.lock() {
|
||||
let mut id = 1;
|
||||
let used: Vec<usize> = st.iter().map(|s| s.id).collect();
|
||||
while used.contains(&id) {
|
||||
id += 1;
|
||||
}
|
||||
id
|
||||
} else {
|
||||
self.next_id
|
||||
};
|
||||
self.next_id = id + 1;
|
||||
|
||||
// Build label from first chars of input
|
||||
let label: String = input.chars().take(10).collect();
|
||||
|
||||
// Add to status
|
||||
if let Ok(mut st) = self.status.lock() {
|
||||
st.push(SessionStatus {
|
||||
id,
|
||||
state: "thinking...".to_string(),
|
||||
label,
|
||||
});
|
||||
}
|
||||
|
||||
// Notify reader thread of the session ID
|
||||
let _ = self.id_tx.send(id);
|
||||
|
||||
pub fn send(&mut self, input: &str) {
|
||||
Self::set_state(&self.state, "thinking...");
|
||||
let msg = serde_json::json!({
|
||||
"type": "user",
|
||||
"message": {
|
||||
"role": "user",
|
||||
"content": input
|
||||
}
|
||||
"message": { "role": "user", "content": input }
|
||||
});
|
||||
|
||||
if let Ok(mut stdin) = self.stdin.lock() {
|
||||
let _ = writeln!(stdin, "{}", msg);
|
||||
let _ = stdin.flush();
|
||||
}
|
||||
|
||||
id
|
||||
let _ = writeln!(self.stdin, "{}", msg);
|
||||
let _ = self.stdin.flush();
|
||||
}
|
||||
|
||||
/// Try to receive one output event. Non-blocking.
|
||||
pub fn try_recv(&self) -> Option<OutputEvent> {
|
||||
self.output_rx.try_recv().ok()
|
||||
}
|
||||
|
||||
/// Cancel all in-progress AI sessions by sending SIGINT to the child process.
|
||||
pub fn cancel(&mut self) {
|
||||
if let Ok(mut st) = self.status.lock() {
|
||||
if st.is_empty() {
|
||||
return;
|
||||
}
|
||||
st.clear();
|
||||
}
|
||||
// Send SIGINT to Claude child process
|
||||
unsafe {
|
||||
libc::kill(self.child_pid as i32, libc::SIGINT);
|
||||
}
|
||||
// Drain any pending events
|
||||
Self::set_state(&self.state, "idle");
|
||||
unsafe { libc::kill(self.child_pid as i32, libc::SIGINT); }
|
||||
while self.output_rx.try_recv().is_ok() {}
|
||||
}
|
||||
|
||||
/// Check if there are any active sessions.
|
||||
pub fn is_busy(&self) -> bool {
|
||||
self.status.lock().map(|st| !st.is_empty()).unwrap_or(false)
|
||||
self.state.lock().map(|s| *s != "idle").unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Get current status string for status bar.
|
||||
pub fn status_line(&self) -> String {
|
||||
let st = match self.status.lock() {
|
||||
Ok(st) => st,
|
||||
Err(_) => return "error".to_string(),
|
||||
};
|
||||
if st.is_empty() {
|
||||
"idle".to_string()
|
||||
} else {
|
||||
st.iter()
|
||||
.map(|s| format!("[{}] {}[{}]", s.id, s.state, s.label))
|
||||
.collect::<Vec<_>>()
|
||||
.join(" | ")
|
||||
}
|
||||
self.state.lock().map(|s| s.clone()).unwrap_or_else(|_| "error".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,24 +2,14 @@ use std::env;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
|
||||
/// Execute a shell command line and print output directly to stdout/stderr.
|
||||
/// Returns the exit code.
|
||||
pub fn execute(input: &str) -> i32 {
|
||||
let trimmed = input.trim();
|
||||
|
||||
// Handle cd specially - it must affect our process
|
||||
if trimmed == "cd" || trimmed.starts_with("cd ") {
|
||||
return handle_cd(trimmed);
|
||||
}
|
||||
|
||||
// Run everything else through sh -c for full shell semantics
|
||||
// (pipes, redirects, globs, etc.)
|
||||
let status = Command::new("sh")
|
||||
.arg("-c")
|
||||
.arg(trimmed)
|
||||
.status();
|
||||
|
||||
match status {
|
||||
match Command::new("sh").arg("-c").arg(trimmed).status() {
|
||||
Ok(s) => s.code().unwrap_or(1),
|
||||
Err(e) => {
|
||||
eprintln!("aishell: {}", e);
|
||||
@@ -28,53 +18,38 @@ pub fn execute(input: &str) -> i32 {
|
||||
}
|
||||
}
|
||||
|
||||
fn home_dir() -> Option<String> {
|
||||
env::var("HOME").ok()
|
||||
}
|
||||
|
||||
fn handle_cd(input: &str) -> i32 {
|
||||
let target = input.strip_prefix("cd").unwrap().trim();
|
||||
|
||||
let dir: PathBuf = if target.is_empty() {
|
||||
// cd with no args → home directory
|
||||
match env::var("HOME") {
|
||||
Ok(home) => PathBuf::from(home),
|
||||
Err(_) => {
|
||||
eprintln!("aishell: cd: HOME not set");
|
||||
return 1;
|
||||
}
|
||||
match home_dir() {
|
||||
Some(h) => PathBuf::from(h),
|
||||
None => { eprintln!("aishell: cd: HOME not set"); return 1; }
|
||||
}
|
||||
} else if target.starts_with('~') {
|
||||
// Expand ~ to HOME
|
||||
match env::var("HOME") {
|
||||
Ok(home) => PathBuf::from(target.replacen('~', &home, 1)),
|
||||
Err(_) => {
|
||||
eprintln!("aishell: cd: HOME not set");
|
||||
return 1;
|
||||
}
|
||||
match home_dir() {
|
||||
Some(h) => PathBuf::from(target.replacen('~', &h, 1)),
|
||||
None => { eprintln!("aishell: cd: HOME not set"); return 1; }
|
||||
}
|
||||
} else if target == "-" {
|
||||
// cd - → previous directory
|
||||
match env::var("OLDPWD") {
|
||||
Ok(old) => {
|
||||
println!("{}", old);
|
||||
PathBuf::from(old)
|
||||
}
|
||||
Err(_) => {
|
||||
eprintln!("aishell: cd: OLDPWD not set");
|
||||
return 1;
|
||||
}
|
||||
Ok(old) => { println!("{}", old); PathBuf::from(old) }
|
||||
Err(_) => { eprintln!("aishell: cd: OLDPWD not set"); return 1; }
|
||||
}
|
||||
} else {
|
||||
PathBuf::from(target)
|
||||
};
|
||||
|
||||
// Save current dir as OLDPWD
|
||||
if let Ok(cwd) = env::current_dir() {
|
||||
env::set_var("OLDPWD", cwd);
|
||||
}
|
||||
|
||||
match env::set_current_dir(&dir) {
|
||||
Ok(_) => 0,
|
||||
Err(e) => {
|
||||
eprintln!("aishell: cd: {}: {}", dir.display(), e);
|
||||
1
|
||||
}
|
||||
Err(e) => { eprintln!("aishell: cd: {}: {}", dir.display(), e); 1 }
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user