diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8a10c170d7..1dac2fcdb4 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -114,8 +114,6 @@ impl UnifiedExecSessionManager { &self, request: UnifiedExecRequest<'_>, ) -> Result { - tracing::error!("In the exec"); - // todo update the errors let timeout_ms = request.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS); let mut new_session: Option = None; @@ -125,9 +123,15 @@ impl UnifiedExecSessionManager { let output_notify; if let Some(existing_id) = request.session_id { - let sessions = self.sessions.lock().await; + let mut sessions = self.sessions.lock().await; match sessions.get(&existing_id) { Some(session) => { + if session.has_exited() { + sessions.remove(&existing_id); + return Err(UnifiedExecError::UnknownSessionId { + session_id: existing_id, + }); + } let (buffer, notify) = session.output_handles(); session_id = existing_id; writer_tx = session.writer_sender(); @@ -140,6 +144,7 @@ impl UnifiedExecSessionManager { }); } } + drop(sessions); } else { let command = command_from_chunks(request.input_chunks)?; let new_id = self.next_session_id.fetch_add(1, Ordering::SeqCst); @@ -203,6 +208,18 @@ impl UnifiedExecSessionManager { let should_store_session = if let Some(session) = new_session.as_ref() { !session.has_exited() + } else if request.session_id.is_some() { + let mut sessions = self.sessions.lock().await; + if let Some(existing) = sessions.get(&session_id) { + if existing.has_exited() { + sessions.remove(&session_id); + false + } else { + true + } + } else { + false + } } else { true }; @@ -322,8 +339,8 @@ async fn create_unified_exec_session( #[cfg(test)] mod tests { - use super::*; use super::path::parse_command_line; + use super::*; #[test] fn parse_command_line_splits_words() { @@ -355,7 +372,7 @@ mod tests { let open_shell = manager .handle_request(UnifiedExecRequest { session_id: None, - input_chunks: &["/bin/bash".to_string(), "-i".to_string()], + input_chunks: &["bash".to_string(), "-i".to_string()], timeout_ms: Some(1_500), }) .await?; @@ -492,13 +509,78 @@ mod tests { Ok(()) } + #[cfg(unix)] + #[tokio::test] + async fn correct_path_resolution() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + let result = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["echo".to_string(), "codex".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + assert!(result.session_id.is_none()); + assert!(result.output.contains("codex")); + + assert!(manager.sessions.lock().await.is_empty()); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reusing_completed_session_returns_unknown_session() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let open_shell = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["/bin/bash".to_string(), "-i".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + let session_id = open_shell.session_id.expect("expected session id"); + + manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &["exit\n".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + tokio::time::sleep(Duration::from_millis(200)).await; + + let err = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &[], + timeout_ms: Some(100), + }) + .await + .expect_err("expected unknown session error"); + + match err { + UnifiedExecError::UnknownSessionId { session_id: err_id } => { + assert_eq!(err_id, session_id); + } + other => panic!("expected UnknownSessionId, got {other:?}"), + } + + assert!(!manager.sessions.lock().await.contains_key(&session_id)); + + Ok(()) + } + #[test] fn truncate_middle_no_newlines_fallback() { - let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"; - let max_bytes = 16; + let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*"; + let max_bytes = 32; let (out, original) = truncate_middle(s, max_bytes); - assert_eq!(out, "…16 tokens truncated…"); - assert_eq!(original, Some(16)); + assert_eq!(out, "abcdef[TRUNCATED CONTENT]UVWXYZ*"); + assert_eq!(original, Some((s.len() as u64).div_ceil(4))); } #[test] @@ -510,12 +592,23 @@ mod tests { assert_eq!(s.len(), 80); let max_bytes = 64; - assert_eq!( - truncate_middle(&s, max_bytes), - ( - "001\n002\n003\n004\n…12 tokens truncated…\n017\n018\n019\n020\n".to_string(), - Some(20) - ) - ); + let (out, tokens) = truncate_middle(&s, max_bytes); + assert!(out.starts_with("001\n002\n003\n004\n")); + assert!(out.contains("[TRUNCATED CONTENT]")); + assert!(out.ends_with("017\n018\n019\n020\n")); + assert_eq!(tokens, Some(20)); + } + + #[test] + // Ensure truncation is resilient against multi-bytes chars split. + fn truncate_middle_handles_utf8_content() { + let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with ascii text\n"; + let max_bytes = 32; + let (out, tokens) = truncate_middle(s, max_bytes); + + assert!(out.contains("[TRUNCATED CONTENT]")); + assert!(out.chars().any(|c| c == '😀')); + assert!(!out.contains('\u{fffd}')); + assert_eq!(tokens, Some((s.len() as u64).div_ceil(4))); } } diff --git a/codex-rs/core/src/unified_exec/truncation.rs b/codex-rs/core/src/unified_exec/truncation.rs index 7889c66d72..540ac563be 100644 --- a/codex-rs/core/src/unified_exec/truncation.rs +++ b/codex-rs/core/src/unified_exec/truncation.rs @@ -1,10 +1,13 @@ +const TRUNCATION_MARKER: &str = "[TRUNCATED CONTENT]"; + pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { if s.len() <= max_bytes { return (s.to_string(), None); } + let est_tokens = (s.len() as u64).div_ceil(4); if max_bytes == 0 { - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); } fn truncate_on_boundary(input: &str, max_len: usize) -> &str { @@ -41,50 +44,28 @@ pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option idx } - let mut guess_tokens = est_tokens; - for _ in 0..4 { - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let mut suffix_start = pick_suffix_start(s, right_budget); - if suffix_start < prefix_end { - suffix_start = prefix_end; - } - let kept_content_bytes = prefix_end + (s.len() - suffix_start); - let truncated_content_bytes = s.len().saturating_sub(kept_content_bytes); - let new_tokens = (truncated_content_bytes as u64).div_ceil(4); - if new_tokens == guess_tokens { - let mut out = String::with_capacity(marker_len + kept_content_bytes + 1); - out.push_str(&s[..prefix_end]); - out.push_str(&marker); - out.push('\n'); - out.push_str(&s[suffix_start..]); - return (out, Some(est_tokens)); - } - guess_tokens = new_tokens; + let marker_len = TRUNCATION_MARKER.len(); + if marker_len >= max_bytes { + return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); } - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); let keep_budget = max_bytes.saturating_sub(marker_len); if keep_budget == 0 { - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + return (TRUNCATION_MARKER.to_string(), Some(est_tokens)); } + let left_budget = keep_budget / 2; let right_budget = keep_budget - left_budget; let prefix_end = pick_prefix_end(s, left_budget); - let suffix_start = pick_suffix_start(s, right_budget); - let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start) + 1); + let mut suffix_start = pick_suffix_start(s, right_budget); + if suffix_start < prefix_end { + suffix_start = prefix_end; + } + + let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start)); out.push_str(&s[..prefix_end]); - out.push_str(&marker); - out.push('\n'); + out.push_str(TRUNCATION_MARKER); out.push_str(&s[suffix_start..]); + (out, Some(est_tokens)) }