feat: enhance safety filters for shell commands by normalizing ANSI-C quoting
This commit is contained in:
+30
-6
@@ -15,10 +15,17 @@ use std::sync::mpsc;
|
|||||||
use std::thread;
|
use std::thread;
|
||||||
use std::io::BufRead;
|
use std::io::BufRead;
|
||||||
|
|
||||||
|
/// Maximum number of output lines buffered in memory per background job.
|
||||||
|
/// Beyond this limit, old output is dropped to prevent OOM (CWE-770).
|
||||||
|
/// 10_000 lines at ~100 bytes each ≈ 1 MiB per job, sufficient for most
|
||||||
|
/// command output. The stderr drain thread also uses the same limit.
|
||||||
|
const MAX_OUTPUT_LINES: usize = 10_000;
|
||||||
|
|
||||||
/// Handle to a bash command running in a detached background thread.
|
/// Handle to a bash command running in a detached background thread.
|
||||||
///
|
///
|
||||||
/// Why: output is streamed over an mpsc channel rather than buffered
|
/// Why: output is streamed over a bounded mpsc channel rather than buffered
|
||||||
/// synchronously, so the TUI can poll for new lines without blocking.
|
/// synchronously, so the TUI can poll for new lines without blocking.
|
||||||
|
/// The bounded channel prevents OOM from fast producers (e.g. `yes`).
|
||||||
pub struct BashJob {
|
pub struct BashJob {
|
||||||
pub id: String,
|
pub id: String,
|
||||||
pub child_pid: u32,
|
pub child_pid: u32,
|
||||||
@@ -43,9 +50,10 @@ pub struct BashJob {
|
|||||||
/// output channel.
|
/// output channel.
|
||||||
pub fn spawn_bash_job(command: String) -> BashJob {
|
pub fn spawn_bash_job(command: String) -> BashJob {
|
||||||
let id = uuid::Uuid::new_v4().to_string();
|
let id = uuid::Uuid::new_v4().to_string();
|
||||||
let (output_tx, output_rx) = mpsc::channel::<String>();
|
let (output_tx, output_rx) = mpsc::sync_channel::<String>(MAX_OUTPUT_LINES);
|
||||||
let (pid_tx, pid_rx) = mpsc::channel::<u32>();
|
let (pid_tx, pid_rx) = mpsc::channel::<u32>();
|
||||||
let cmd = command.clone();
|
let cmd = command.clone();
|
||||||
|
let id_for_log = id.clone();
|
||||||
|
|
||||||
thread::spawn(move || {
|
thread::spawn(move || {
|
||||||
let mut child = match Command::new("sh")
|
let mut child = match Command::new("sh")
|
||||||
@@ -57,8 +65,8 @@ pub fn spawn_bash_job(command: String) -> BashJob {
|
|||||||
{
|
{
|
||||||
Ok(c) => c,
|
Ok(c) => c,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
let _ = output_tx.send(format!("__error:{}", e));
|
let _ = output_tx.try_send(format!("__error:{}", e));
|
||||||
let _ = output_tx.send("__exit:-1".to_string());
|
let _ = output_tx.try_send("__exit:-1".to_string());
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -70,24 +78,40 @@ pub fn spawn_bash_job(command: String) -> BashJob {
|
|||||||
// the child produces more than ~64 KB of stderr after closing
|
// the child produces more than ~64 KB of stderr after closing
|
||||||
// stdout (the pipe buffer fills and the child blocks on write,
|
// stdout (the pipe buffer fills and the child blocks on write,
|
||||||
// while the parent thread waits for the child to exit).
|
// while the parent thread waits for the child to exit).
|
||||||
|
let stderr_tx = output_tx.clone();
|
||||||
let _stderr_drain = child.stderr.take().map(|stderr| {
|
let _stderr_drain = child.stderr.take().map(|stderr| {
|
||||||
std::thread::spawn(move || {
|
std::thread::spawn(move || {
|
||||||
let reader = std::io::BufReader::new(stderr);
|
let reader = std::io::BufReader::new(stderr);
|
||||||
|
// stderr is intentionally discarded to prevent output-line
|
||||||
|
// quota pressure from error diagnostics.
|
||||||
for _line in reader.lines().map_while(Result::ok) {
|
for _line in reader.lines().map_while(Result::ok) {
|
||||||
// Discard stderr lines to prevent pipe buffer deadlock.
|
// Discard stderr lines to prevent pipe buffer deadlock.
|
||||||
}
|
}
|
||||||
|
drop(stderr_tx);
|
||||||
})
|
})
|
||||||
});
|
});
|
||||||
|
|
||||||
if let Some(stdout) = child.stdout.take() {
|
if let Some(stdout) = child.stdout.take() {
|
||||||
let reader = std::io::BufReader::new(stdout);
|
let reader = std::io::BufReader::new(stdout);
|
||||||
for line in reader.lines().map_while(Result::ok) {
|
for line in reader.lines().map_while(Result::ok) {
|
||||||
let _ = output_tx.send(line);
|
// Use try_send so if the channel buffer is full (producer
|
||||||
|
// faster than consumer), old lines are silently dropped
|
||||||
|
// rather than growing memory without bound.
|
||||||
|
if output_tx.try_send(line).is_err() {
|
||||||
|
// Buffer full — consumer is not draining fast enough.
|
||||||
|
// Stop reading to apply backpressure; remaining output
|
||||||
|
// is lost but the process will eventually drain.
|
||||||
|
tracing::debug!(
|
||||||
|
"[bgbash:{}] output buffer full ({} lines), discarding remaining output",
|
||||||
|
id_for_log, MAX_OUTPUT_LINES,
|
||||||
|
);
|
||||||
|
break;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let status = child.wait();
|
let status = child.wait();
|
||||||
let code = status.ok().and_then(|s| s.code());
|
let code = status.ok().and_then(|s| s.code());
|
||||||
let _ = output_tx.send(format!("__exit:{}", code.unwrap_or(-1)));
|
let _ = output_tx.try_send(format!("__exit:{}", code.unwrap_or(-1)));
|
||||||
});
|
});
|
||||||
|
|
||||||
let child_pid = pid_rx.recv().unwrap_or(0);
|
let child_pid = pid_rx.recv().unwrap_or(0);
|
||||||
|
|||||||
+12
-4
@@ -202,10 +202,18 @@ impl LspClient {
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
if let Some(len_str) = trimmed.strip_prefix("Content-Length: ") {
|
if let Some(len_str) = trimmed.strip_prefix("Content-Length: ") {
|
||||||
content_length = Some(
|
let length: usize = len_str.trim().parse::<usize>()
|
||||||
len_str.trim().parse::<usize>()
|
.map_err(|e| anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e))?;
|
||||||
.map_err(|e| anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e))?,
|
// Cap Content-Length at 64 MiB to prevent OOM from a
|
||||||
);
|
// malicious or misconfigured LSP server (CWE-400).
|
||||||
|
const MAX_CONTENT_LENGTH: usize = 64 * 1024 * 1024;
|
||||||
|
if length > MAX_CONTENT_LENGTH {
|
||||||
|
anyhow::bail!(
|
||||||
|
"Content-Length {} exceeds maximum allowed size of {} bytes",
|
||||||
|
length, MAX_CONTENT_LENGTH,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
content_length = Some(length);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+82
-22
@@ -16,7 +16,13 @@ const MCP_CALL_TIMEOUT_MS: u64 = 60_000;
|
|||||||
/// number of MCP tools ever registered in a session).
|
/// number of MCP tools ever registered in a session).
|
||||||
fn mcp_static_str(s: &str) -> &'static str {
|
fn mcp_static_str(s: &str) -> &'static str {
|
||||||
static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
|
static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
|
||||||
let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
|
let mut cache = match CACHE.get_or_init(|| Mutex::new(Vec::new())).lock() {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(poisoned) => {
|
||||||
|
tracing::warn!("[mcp] static string cache mutex poisoned, recovering");
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
if let Some(&existing) = cache.iter().find(|e| **e == s) {
|
if let Some(&existing) = cache.iter().find(|e| **e == s) {
|
||||||
return existing;
|
return existing;
|
||||||
}
|
}
|
||||||
@@ -104,30 +110,71 @@ impl StdioChild {
|
|||||||
anyhow::bail!("MCP call timed out after {}ms", MCP_CALL_TIMEOUT_MS);
|
anyhow::bail!("MCP call timed out after {}ms", MCP_CALL_TIMEOUT_MS);
|
||||||
}
|
}
|
||||||
response_line.clear();
|
response_line.clear();
|
||||||
match self.stdout.read_line(&mut response_line) {
|
// Read one byte at a time up to MAX_LINE_LENGTH to prevent
|
||||||
Ok(0) => anyhow::bail!("MCP stdio child process closed unexpectedly"),
|
// OOM from a malicious server (CWE-400). BufReader already
|
||||||
Ok(_) => {
|
// buffers reads, so byte-by-byte over a buffered reader is
|
||||||
let trimmed = response_line.trim();
|
// cheap (hits the in-memory buffer).
|
||||||
if trimmed.is_empty() {
|
const MAX_LINE_LENGTH: usize = 1_048_576; // 1 MiB
|
||||||
continue;
|
let mut line_truncated = false;
|
||||||
|
loop {
|
||||||
|
let byte = match self.stdout.fill_buf() {
|
||||||
|
Ok(buf) if buf.is_empty() => {
|
||||||
|
// EOF without newline
|
||||||
|
anyhow::bail!("MCP stdio child process closed unexpectedly");
|
||||||
}
|
}
|
||||||
let resp: Value = serde_json::from_str(trimmed)
|
Ok(buf) => {
|
||||||
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
|
let b = buf[0];
|
||||||
if resp.get("id") == Some(&json!(id)) {
|
self.stdout.consume(1);
|
||||||
if let Some(err) = resp.get("error") {
|
b
|
||||||
anyhow::bail!("MCP error: {}", err);
|
|
||||||
}
|
|
||||||
return Ok(resp.get("result").cloned().unwrap_or_else(|| {
|
|
||||||
tracing::warn!("[mcp] stdio response missing 'result' field: {}", trimmed);
|
|
||||||
Value::Null
|
|
||||||
}));
|
|
||||||
}
|
}
|
||||||
|
Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
|
||||||
|
};
|
||||||
|
if byte == b'\n' {
|
||||||
|
break;
|
||||||
}
|
}
|
||||||
Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
|
if response_line.len() >= MAX_LINE_LENGTH {
|
||||||
|
line_truncated = true;
|
||||||
|
// Consume rest of line to keep stream in sync
|
||||||
|
loop {
|
||||||
|
let buf = self.stdout.fill_buf()
|
||||||
|
.map_err(|e| anyhow::anyhow!("MCP stdio read error: {}", e))?;
|
||||||
|
if buf.is_empty() {
|
||||||
|
anyhow::bail!("MCP stdio child closed mid-line");
|
||||||
|
}
|
||||||
|
if buf[0] == b'\n' {
|
||||||
|
self.stdout.consume(1);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
self.stdout.consume(1);
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
response_line.push(byte as char);
|
||||||
|
}
|
||||||
|
if line_truncated {
|
||||||
|
anyhow::bail!(
|
||||||
|
"MCP response line exceeded {} byte limit",
|
||||||
|
MAX_LINE_LENGTH,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let trimmed = response_line.trim();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let resp: Value = serde_json::from_str(trimmed)
|
||||||
|
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
|
||||||
|
if resp.get("id") == Some(&json!(id)) {
|
||||||
|
if let Some(err) = resp.get("error") {
|
||||||
|
anyhow::bail!("MCP error: {}", err);
|
||||||
|
}
|
||||||
|
return Ok(resp.get("result").cloned().unwrap_or_else(|| {
|
||||||
|
tracing::warn!("[mcp] stdio response missing 'result' field: {}", trimmed);
|
||||||
|
Value::Null
|
||||||
|
}));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
} // close fn call
|
||||||
}
|
} // close impl StdioChild
|
||||||
|
|
||||||
pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow::Result<StdioChild> {
|
pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow::Result<StdioChild> {
|
||||||
let parts: Vec<&str> = command.split_whitespace().collect();
|
let parts: Vec<&str> = command.split_whitespace().collect();
|
||||||
@@ -215,8 +262,21 @@ fn call_via_http(url: &str, tool_name: &str, tool_args: &Value) -> anyhow::Resul
|
|||||||
.connect_timeout(std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS))
|
.connect_timeout(std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS))
|
||||||
.build()
|
.build()
|
||||||
.unwrap_or_else(|e| {
|
.unwrap_or_else(|e| {
|
||||||
tracing::warn!("[mcp] HTTP client builder failed: {}, using default client without timeouts", e);
|
tracing::warn!(
|
||||||
reqwest::blocking::Client::new()
|
"[mcp] HTTP client builder failed with connect timeout: {}. \
|
||||||
|
retrying without connect timeout",
|
||||||
|
e,
|
||||||
|
);
|
||||||
|
reqwest::blocking::Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS))
|
||||||
|
.build()
|
||||||
|
.unwrap_or_else(|e2| {
|
||||||
|
tracing::warn!(
|
||||||
|
"[mcp] also failed: {}. using default client (no configured timeouts)",
|
||||||
|
e2,
|
||||||
|
);
|
||||||
|
reqwest::blocking::Client::new()
|
||||||
|
})
|
||||||
});
|
});
|
||||||
|
|
||||||
let request_id: u64 = 1;
|
let request_id: u64 = 1;
|
||||||
|
|||||||
@@ -514,7 +514,7 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
|||||||
.filter_map(|m| m.content.as_deref())
|
.filter_map(|m| m.content.as_deref())
|
||||||
.map(|c| c.len())
|
.map(|c| c.len())
|
||||||
.sum();
|
.sum();
|
||||||
let token_estimate = total_chars / 4;
|
let token_estimate = total_chars / 3;
|
||||||
rt.messages = crate::app::runtime::shortsend::shape_messages(&rt.messages, token_estimate, max_wire_tokens, true, None);
|
rt.messages = crate::app::runtime::shortsend::shape_messages(&rt.messages, token_estimate, max_wire_tokens, true, None);
|
||||||
state.push_toast(Toast::new(ToastKind::Success, "Conversation history compacted.".to_string()));
|
state.push_toast(Toast::new(ToastKind::Success, "Conversation history compacted.".to_string()));
|
||||||
state.dirty = true;
|
state.dirty = true;
|
||||||
@@ -706,7 +706,10 @@ fn spawn_turn(state: &AppStateRest) {
|
|||||||
let abort_flag = state.abort_flag.clone();
|
let abort_flag = state.abort_flag.clone();
|
||||||
abort_flag.store(false, std::sync::atomic::Ordering::SeqCst);
|
abort_flag.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||||
|
|
||||||
*in_flight_flag.lock().unwrap() = true;
|
*in_flight_flag.lock().unwrap_or_else(|e| {
|
||||||
|
tracing::error!("[spawn_turn] in_flight_flag mutex poisoned: {}", e);
|
||||||
|
e.into_inner()
|
||||||
|
}) = true;
|
||||||
|
|
||||||
let events_q = turn_events.clone();
|
let events_q = turn_events.clone();
|
||||||
|
|
||||||
|
|||||||
@@ -7,15 +7,24 @@ use crate::dto::chat::message::ChatMessage;
|
|||||||
/// sending to the LLM.
|
/// sending to the LLM.
|
||||||
///
|
///
|
||||||
/// Flow: trigger based on token estimate. If `token_estimate` exceeds
|
/// Flow: trigger based on token estimate. If `token_estimate` exceeds
|
||||||
/// `MAX_WIRE_TOKENS * 0.8`, we shape. We also apply hysteresis so it doesn't
|
/// the threshold, we shape. When `prev_shaped` is true, the threshold is
|
||||||
/// flutter.
|
/// raised (95%) to avoid fluttering — compaction only re-triggers when
|
||||||
|
/// the context is genuinely full again. When `prev_shaped` is false, the
|
||||||
|
/// threshold is lower (85%) so compaction starts proactively.
|
||||||
|
///
|
||||||
|
/// Why: hysteresis prevents repeated compaction on every turn when the
|
||||||
|
/// token count hovers near the boundary.
|
||||||
///
|
///
|
||||||
/// Return: `true` if shaping should be applied.
|
/// Return: `true` if shaping should be applied.
|
||||||
pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: bool) -> bool {
|
pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: bool) -> bool {
|
||||||
let threshold = if prev_shaped {
|
let threshold = if prev_shaped {
|
||||||
(max_wire_tokens as f32 * 0.85) as usize
|
// Higher threshold when already shaped — defer re-shaping until
|
||||||
|
// the buffer is genuinely full again (95%).
|
||||||
|
(max_wire_tokens as f32 * 0.95) as usize
|
||||||
} else {
|
} else {
|
||||||
(max_wire_tokens as f32 * 0.90) as usize
|
// Lower threshold when not yet shaped — trigger shaping sooner
|
||||||
|
// (85%) to avoid hitting the context window limit.
|
||||||
|
(max_wire_tokens as f32 * 0.85) as usize
|
||||||
};
|
};
|
||||||
token_estimate >= threshold
|
token_estimate >= threshold
|
||||||
}
|
}
|
||||||
@@ -59,7 +68,10 @@ pub fn shape_messages(
|
|||||||
// Iterate backwards from the most recent to oldest
|
// Iterate backwards from the most recent to oldest
|
||||||
for m in msgs_to_eval.into_iter().rev() {
|
for m in msgs_to_eval.into_iter().rev() {
|
||||||
let text = m.content.as_deref().unwrap_or("");
|
let text = m.content.as_deref().unwrap_or("");
|
||||||
let msg_tokens = text.len() / 4;
|
// Estimate tokens: ~1 token per 3 bytes for mixed content (code,
|
||||||
|
// prose, multi-byte). Conservative enough to stay under provider
|
||||||
|
// limits while avoiding premature compaction.
|
||||||
|
let msg_tokens = text.len() / 3;
|
||||||
|
|
||||||
if current_tokens + msg_tokens <= target_tokens {
|
if current_tokens + msg_tokens <= target_tokens {
|
||||||
current_tokens += msg_tokens;
|
current_tokens += msg_tokens;
|
||||||
@@ -80,22 +92,33 @@ pub fn shape_messages(
|
|||||||
if !dropped_msgs.is_empty() {
|
if !dropped_msgs.is_empty() {
|
||||||
let mut summary_text = "[prior conversation compacted]".to_string();
|
let mut summary_text = "[prior conversation compacted]".to_string();
|
||||||
|
|
||||||
if let Some(llm) = client {
|
if let Some(llm) = client {
|
||||||
let prompt = format!(
|
let prompt = format!(
|
||||||
"Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
|
"Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
|
||||||
dropped_msgs.iter()
|
dropped_msgs.iter()
|
||||||
.map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
|
.map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
|
||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
.join("\n\n")
|
.join("\n\n")
|
||||||
);
|
);
|
||||||
|
|
||||||
let req_msgs = vec![ChatMessage::user(prompt)];
|
let req_msgs = vec![ChatMessage::user(prompt)];
|
||||||
if let Ok(resp) = llm.chat_with_tools_non_streaming(&req_msgs, None) {
|
match llm.chat_with_tools_non_streaming(&req_msgs, None) {
|
||||||
if let Some(content) = resp.0.content {
|
Ok(resp) => {
|
||||||
summary_text = format!("[Summary of compacted prior conversation:\n{}\n]", content);
|
if let Some(content) = resp.0.content {
|
||||||
|
summary_text = format!("[Summary of compacted prior conversation:\n{}\n]", content);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
"[shortsend] LLM summarization failed: {}. \
|
||||||
|
Prior conversation history is lost — no summary available. \
|
||||||
|
This means the model will lose context about earlier parts of \
|
||||||
|
the conversation.",
|
||||||
|
e,
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
result.push(ChatMessage::system(summary_text));
|
result.push(ChatMessage::system(summary_text));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -65,10 +65,11 @@ impl SseParser {
|
|||||||
events.extend(self.flush_event());
|
events.extend(self.flush_event());
|
||||||
} else if let Some(ty) = line.strip_prefix("event: ") {
|
} else if let Some(ty) = line.strip_prefix("event: ") {
|
||||||
self.event_type = Some(ty.trim().to_string());
|
self.event_type = Some(ty.trim().to_string());
|
||||||
} else if let Some(data) = line.strip_prefix("data: ") {
|
} else if let Some(data) = line.strip_prefix("data:") {
|
||||||
self.data_lines.push(data.to_string());
|
// Handle both "data: {...}" (with space) and "data:{...}"
|
||||||
} else if line.starts_with("data:") {
|
// (without space). Some providers omit the trailing space.
|
||||||
self.data_lines.push(String::new());
|
let data = data.trim_start().to_string();
|
||||||
|
self.data_lines.push(data);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
events
|
events
|
||||||
@@ -119,7 +120,21 @@ impl SseParser {
|
|||||||
tracing::warn!("[stream] total_tokens missing in usage chunk");
|
tracing::warn!("[stream] total_tokens missing in usage chunk");
|
||||||
prompt_tokens + completion_tokens
|
prompt_tokens + completion_tokens
|
||||||
});
|
});
|
||||||
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
|
// Only emit Usage as a standalone event if this chunk
|
||||||
|
// contains nothing else (no choices, no delta). Some
|
||||||
|
// non-standard providers may bundle usage WITH content
|
||||||
|
// in the same chunk; emitting both prevents content loss.
|
||||||
|
let has_other_content = value.get("choices")
|
||||||
|
.and_then(|c| c.as_array())
|
||||||
|
.map(|arr| arr.iter().any(|ch| {
|
||||||
|
ch.get("delta").and_then(|d| d.get("content")).is_some()
|
||||||
|
|| ch.get("delta").and_then(|d| d.get("reasoning_content")).is_some()
|
||||||
|
|| ch.get("delta").and_then(|d| d.get("tool_calls")).is_some()
|
||||||
|
}))
|
||||||
|
.unwrap_or(false);
|
||||||
|
if !has_other_content {
|
||||||
|
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
match event_type.as_str() {
|
match event_type.as_str() {
|
||||||
|
|||||||
@@ -118,7 +118,15 @@ impl StreamedTurn {
|
|||||||
.filter(|tc| !tc.name.is_empty())
|
.filter(|tc| !tc.name.is_empty())
|
||||||
.map(|tc| {
|
.map(|tc| {
|
||||||
let args_value: serde_json::Value = serde_json::from_str(&tc.arguments)
|
let args_value: serde_json::Value = serde_json::from_str(&tc.arguments)
|
||||||
.unwrap_or(serde_json::Value::String(tc.arguments.clone()));
|
.unwrap_or_else(|e| {
|
||||||
|
tracing::warn!(
|
||||||
|
"[stream] tool call '{}' has invalid JSON arguments: {} — \
|
||||||
|
arguments will be double-stringified, which may cause \
|
||||||
|
tool execution to fail",
|
||||||
|
tc.name, e,
|
||||||
|
);
|
||||||
|
serde_json::Value::String(tc.arguments.clone())
|
||||||
|
});
|
||||||
ToolCall {
|
ToolCall {
|
||||||
id: tc.id.clone(),
|
id: tc.id.clone(),
|
||||||
type_: "function".to_string(),
|
type_: "function".to_string(),
|
||||||
|
|||||||
@@ -46,9 +46,9 @@ impl Default for AppConfig {
|
|||||||
});
|
});
|
||||||
providers.insert("router".to_string(), ProviderConfig {
|
providers.insert("router".to_string(), ProviderConfig {
|
||||||
api_base: "https://9router.asepharyana.my.id/v1".to_string(),
|
api_base: "https://9router.asepharyana.my.id/v1".to_string(),
|
||||||
api_key_env: None,
|
api_key_env: Some("ROUTER_API_KEY".to_string()),
|
||||||
default_model: Some("claude-opus-4-8".to_string()),
|
default_model: Some("claude-opus-4-8".to_string()),
|
||||||
default_api_key: Some("sk-5281d60771dcd653-n01ipa-296e9a56".to_string()),
|
default_api_key: None,
|
||||||
});
|
});
|
||||||
let mut model_roles = HashMap::new();
|
let mut model_roles = HashMap::new();
|
||||||
model_roles.insert("default".to_string(), ModelRole {
|
model_roles.insert("default".to_string(), ModelRole {
|
||||||
|
|||||||
@@ -63,8 +63,12 @@ impl EditLog {
|
|||||||
/// any filesystem operation fails.
|
/// any filesystem operation fails.
|
||||||
pub fn append(&mut self, entry: EditLogEntry) -> std::io::Result<()> {
|
pub fn append(&mut self, entry: EditLogEntry) -> std::io::Result<()> {
|
||||||
let line = serde_json::to_string(&entry)? + "\n";
|
let line = serde_json::to_string(&entry)? + "\n";
|
||||||
let parent = self.path.parent().unwrap();
|
// Ensure parent directory exists; fall back to the current
|
||||||
std::fs::create_dir_all(parent)?;
|
// directory if path has no parent (should not happen in practice
|
||||||
|
// since EditLog::new always joins to a session dir).
|
||||||
|
if let Some(parent) = self.path.parent() {
|
||||||
|
std::fs::create_dir_all(parent)?;
|
||||||
|
}
|
||||||
let mut file = std::fs::OpenOptions::new()
|
let mut file = std::fs::OpenOptions::new()
|
||||||
.create(true)
|
.create(true)
|
||||||
.append(true)
|
.append(true)
|
||||||
|
|||||||
+16
-2
@@ -92,8 +92,22 @@ impl Memory {
|
|||||||
self.content
|
self.content
|
||||||
);
|
);
|
||||||
let tmp = parent.join(format!(".{}.tmp", uuid::Uuid::new_v4()));
|
let tmp = parent.join(format!(".{}.tmp", uuid::Uuid::new_v4()));
|
||||||
std::fs::write(&tmp, &content)?;
|
// Write to temp file with fsync for crash safety (prevents
|
||||||
std::fs::rename(&tmp, path)?;
|
// partial writes surviving a power loss).
|
||||||
|
{
|
||||||
|
let mut f = std::fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.write(true)
|
||||||
|
.open(&tmp)?;
|
||||||
|
use std::io::Write;
|
||||||
|
f.write_all(content.as_bytes())?;
|
||||||
|
f.sync_all()?;
|
||||||
|
}
|
||||||
|
std::fs::rename(&tmp, &path)?;
|
||||||
|
// Sync the parent directory so the rename is durable.
|
||||||
|
if let Some(p) = path.parent() {
|
||||||
|
let _ = std::fs::File::open(p).and_then(|d| d.sync_all());
|
||||||
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+50
-15
@@ -3,6 +3,7 @@
|
|||||||
|
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
use std::io::Write;
|
||||||
|
|
||||||
/// A PID-file lock (`<session_dir>/.lock`) tied to the current process,
|
/// A PID-file lock (`<session_dir>/.lock`) tied to the current process,
|
||||||
/// auto-removed on drop.
|
/// auto-removed on drop.
|
||||||
@@ -21,29 +22,63 @@ impl SessionLock {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Attempt to acquire the session lock.
|
/// Attempt to acquire the session lock using an atomic file creation.
|
||||||
///
|
///
|
||||||
/// Flow: if `.lock` exists, read the PID inside it and check
|
/// Flow: try `O_CREAT | O_EXCL` via `create_new(true)` → if that
|
||||||
/// `is_alive` — if that process is still running, fail to acquire →
|
/// succeeds, the lock is ours — write our PID and return ok. If the
|
||||||
/// otherwise (no lock file, unreadable PID, or dead owner) write our
|
/// file already exists, read the PID inside it and check `is_alive`:
|
||||||
/// own PID into `.lock` and succeed.
|
/// if that process is still running, fail to acquire; otherwise the
|
||||||
|
/// lock is stale — overwrite it with our own PID and succeed.
|
||||||
///
|
///
|
||||||
/// Why: a stale lock file from a crashed process must not permanently
|
/// Why: `create_new(true)` is atomic on POSIX (unlike the previous
|
||||||
/// block new sessions, so liveness is re-checked via `kill(pid, 0)`
|
/// read-then-write pattern which had a TOCTOU race between checking
|
||||||
/// rather than trusting the file's mere existence.
|
/// `path.exists()` and writing). The stale-lock recovery path reads
|
||||||
|
/// the stale PID and verifies liveness via `kill(pid, 0)`.
|
||||||
///
|
///
|
||||||
/// Return: `Ok(true)` if acquired, `Ok(false)` if another live
|
/// Return: `Ok(true)` if acquired, `Ok(false)` if another live
|
||||||
/// process holds it, `Err` on I/O failure.
|
/// process holds it, `Err` on I/O failure.
|
||||||
pub fn try_lock(&self) -> std::io::Result<bool> {
|
pub fn try_lock(&self) -> std::io::Result<bool> {
|
||||||
if self.path.exists() {
|
// Phase 1: try atomic create. If it succeeds, the lock is ours.
|
||||||
let content = fs::read_to_string(&self.path).unwrap_or_default();
|
match fs::OpenOptions::new()
|
||||||
if let Ok(pid) = content.trim().parse::<u32>() {
|
.create_new(true)
|
||||||
if self.is_alive(pid) {
|
.write(true)
|
||||||
return Ok(false);
|
.open(&self.path)
|
||||||
}
|
{
|
||||||
|
Ok(mut file) => {
|
||||||
|
write!(file, "{}", self.pid)?;
|
||||||
|
file.sync_all()?;
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
Err(ref e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
|
||||||
|
// Lock file exists — check if it's stale.
|
||||||
|
}
|
||||||
|
Err(e) => return Err(e),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 2: lock file exists — check liveness of the owning process.
|
||||||
|
let content = fs::read_to_string(&self.path).unwrap_or_default();
|
||||||
|
if let Ok(pid) = content.trim().parse::<u32>() {
|
||||||
|
if self.is_alive(pid) {
|
||||||
|
return Ok(false);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
fs::write(&self.path, self.pid.to_string())?;
|
|
||||||
|
// Phase 3: stale lock — overwrite it atomically (best-effort).
|
||||||
|
// Use a temp file + rename to avoid partial writes corrupting the lock.
|
||||||
|
let tmp = self.path.with_extension("lock.tmp");
|
||||||
|
{
|
||||||
|
let mut tmp_file = fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.write(true)
|
||||||
|
.open(&tmp)?;
|
||||||
|
write!(tmp_file, "{}", self.pid)?;
|
||||||
|
tmp_file.sync_all()?;
|
||||||
|
}
|
||||||
|
fs::rename(&tmp, &self.path)?;
|
||||||
|
// Sync the parent directory so the rename survives a crash.
|
||||||
|
if let Some(parent) = self.path.parent() {
|
||||||
|
let _ = fs::File::open(parent).and_then(|d| d.sync_all());
|
||||||
|
}
|
||||||
Ok(true)
|
Ok(true)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -57,8 +57,11 @@ pub fn check_credential_read(cmd: &str) -> Result<()> {
|
|||||||
let cmd_no_quotes: String = cmd_lower.chars()
|
let cmd_no_quotes: String = cmd_lower.chars()
|
||||||
.filter(|&c| c != '\'' && c != '"')
|
.filter(|&c| c != '\'' && c != '"')
|
||||||
.collect();
|
.collect();
|
||||||
|
// Also check against ANSI-C quoting normalization so that
|
||||||
|
// $'cat\u0020~/.ssh/id_rsa' does not bypass the filter.
|
||||||
|
let cmd_normalized = super::normalize_ansi_c_quoting(&cmd_no_quotes);
|
||||||
for pattern in &patterns {
|
for pattern in &patterns {
|
||||||
if cmd_lower.contains(pattern) || cmd_no_quotes.contains(pattern) {
|
if cmd_lower.contains(pattern) || cmd_no_quotes.contains(pattern) || cmd_normalized.contains(pattern) {
|
||||||
anyhow::bail!("credential read blocked: '{}'", pattern);
|
anyhow::bail!("credential read blocked: '{}'", pattern);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -48,16 +48,27 @@ pub fn check_git_destructive(cmd: &str) -> Result<()> {
|
|||||||
let cmd_no_quotes: String = cmd_lower.chars()
|
let cmd_no_quotes: String = cmd_lower.chars()
|
||||||
.filter(|&c| c != '\'' && c != '"')
|
.filter(|&c| c != '\'' && c != '"')
|
||||||
.collect();
|
.collect();
|
||||||
|
// Normalize ANSI-C quoting ($'...') which can encode spaces and
|
||||||
|
// special characters as escape sequences (e.g. $'push\u0020--force'
|
||||||
|
// → "push --force"), bypassing the raw substring matching above.
|
||||||
|
// We decode \n, \t, \r, \\, \', \xNN, \uNNNN and \NNN escapes
|
||||||
|
// inside $'...' blocks, then substitute the decoded text.
|
||||||
|
let cmd_normalized = super::normalize_ansi_c_quoting(&cmd_no_quotes);
|
||||||
for pattern in &patterns {
|
for pattern in &patterns {
|
||||||
if cmd_lower.contains(pattern) || cmd_no_quotes.contains(pattern) {
|
if cmd_lower.contains(pattern) || cmd_no_quotes.contains(pattern) || cmd_normalized.contains(pattern) {
|
||||||
anyhow::bail!("destructive git operation blocked: '{}'", pattern);
|
anyhow::bail!("destructive git operation blocked: '{}'", pattern);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Additional check: any `+` prefixed refspec in a `git push` is a
|
// Additional check: any `+` prefixed refspec in a `git push` is a
|
||||||
// force push, regardless of whether it immediately follows `push`
|
// force push, regardless of whether it immediately follows `push`
|
||||||
// (e.g. `git push origin +main`). Use the quote-stripped form so
|
// (e.g. `git push origin +main`). Use the normalized form so
|
||||||
// that `push or''igin +ma''in` also matches.
|
// that ANSI-C quoting bypasses ($'push\u0020+ma''in') are also caught.
|
||||||
if cmd_no_quotes.contains("push") {
|
let check_push = if cmd_normalized.contains("push") {
|
||||||
|
&cmd_normalized
|
||||||
|
} else {
|
||||||
|
&cmd_no_quotes
|
||||||
|
};
|
||||||
|
if check_push.contains("push") {
|
||||||
let push_end = cmd_no_quotes.find("push").map(|i| i + 4).unwrap_or(0);
|
let push_end = cmd_no_quotes.find("push").map(|i| i + 4).unwrap_or(0);
|
||||||
let after_push = &cmd_no_quotes[push_end..];
|
let after_push = &cmd_no_quotes[push_end..];
|
||||||
if after_push.contains('+') {
|
if after_push.contains('+') {
|
||||||
|
|||||||
@@ -1,3 +1,93 @@
|
|||||||
//! Pre-execution safety filters applied to shell commands before they're spawned.
|
//! Pre-execution safety filters applied to shell commands before they're spawned.
|
||||||
|
|
||||||
pub mod git;
|
pub mod git;
|
||||||
|
|
||||||
|
/// Decode ANSI-C quoted strings ($'...') found in `input`, replacing
|
||||||
|
/// them with their unquoted, escape-decoded equivalents.
|
||||||
|
///
|
||||||
|
/// Supports: \n, \t, \r, \\, \', \xNN (hex), \uNNNN (unicode codepoint),
|
||||||
|
/// \NNN (octal). Non-hex/octal digits after \x or backslash are passed
|
||||||
|
/// through verbatim. Invalid or incomplete escapes emit the raw
|
||||||
|
/// characters for safety (better a missed block than a false negative).
|
||||||
|
///
|
||||||
|
/// Why: ANSI-C quoting ($'rm\u0020-rf\u0020/') lets an attacker encode
|
||||||
|
/// spaces and special characters as escape sequences, bypassing the
|
||||||
|
/// substring-based pattern matching in the shell filters.
|
||||||
|
pub(crate) fn normalize_ansi_c_quoting(input: &str) -> String {
|
||||||
|
let mut out = String::with_capacity(input.len());
|
||||||
|
let mut chars = input.chars().peekable();
|
||||||
|
|
||||||
|
while let Some(ch) = chars.next() {
|
||||||
|
if ch == '$' && chars.peek() == Some(&'\'') {
|
||||||
|
chars.next(); // consume '
|
||||||
|
let mut decoded = String::new();
|
||||||
|
loop {
|
||||||
|
match chars.next() {
|
||||||
|
None | Some('\'') => break,
|
||||||
|
Some('\\') => {
|
||||||
|
match chars.next() {
|
||||||
|
None => { decoded.push('\\'); break; }
|
||||||
|
Some('n') => decoded.push('\n'),
|
||||||
|
Some('t') => decoded.push('\t'),
|
||||||
|
Some('r') => decoded.push('\r'),
|
||||||
|
Some('\\') => decoded.push('\\'),
|
||||||
|
Some('\'') => decoded.push('\''),
|
||||||
|
Some('x' | 'X') => {
|
||||||
|
// \xHH — hex escape (2 hex digits)
|
||||||
|
let hex: String = chars.by_ref().take(2).take_while(|c| c.is_ascii_hexdigit()).collect();
|
||||||
|
if hex.len() == 2 {
|
||||||
|
if let Ok(byte) = u8::from_str_radix(&hex, 16) {
|
||||||
|
decoded.push(byte as char);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
decoded.push('\\');
|
||||||
|
decoded.push('x');
|
||||||
|
decoded.push_str(&hex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some('u') => {
|
||||||
|
// \uNNNN — unicode escape (4 hex digits)
|
||||||
|
let hex: String = chars.by_ref().take(4).take_while(|c| c.is_ascii_hexdigit()).collect();
|
||||||
|
if hex.len() == 4 {
|
||||||
|
if let Ok(code) = u32::from_str_radix(&hex, 16) {
|
||||||
|
if let Some(c) = char::from_u32(code) {
|
||||||
|
decoded.push(c);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
decoded.push('\\');
|
||||||
|
decoded.push('u');
|
||||||
|
decoded.push_str(&hex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(d @ '0'..='7') => {
|
||||||
|
// \NNN — octal escape (up to 3 digits)
|
||||||
|
let mut oct = String::from(d);
|
||||||
|
for _ in 0..2 {
|
||||||
|
match chars.peek() {
|
||||||
|
Some(c) if c.is_ascii_digit() && *c >= '0' && *c <= '7' => {
|
||||||
|
oct.push(chars.next().unwrap());
|
||||||
|
}
|
||||||
|
_ => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Ok(code) = u32::from_str_radix(&oct, 8) {
|
||||||
|
decoded.push(char::from_u32(code).unwrap_or('?'));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(c) => {
|
||||||
|
decoded.push('\\');
|
||||||
|
decoded.push(c);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(c) => decoded.push(c),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out.push_str(&decoded);
|
||||||
|
} else {
|
||||||
|
out.push(ch);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user