chore: rework unified exec events (#7775)

codex-rs/core/src/unified_exec/async_watcher.rs

@@ -0,0 +1,180 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use tokio::sync::Mutex;
+use tokio::time::Duration;
+use tokio::time::Instant;
+
+use crate::codex::Session;
+use crate::codex::TurnContext;
+use crate::exec::ExecToolCallOutput;
+use crate::exec::StreamOutput;
+use crate::protocol::EventMsg;
+use crate::protocol::ExecCommandOutputDeltaEvent;
+use crate::protocol::ExecCommandSource;
+use crate::protocol::ExecOutputStream;
+use crate::tools::events::ToolEmitter;
+use crate::tools::events::ToolEventCtx;
+use crate::tools::events::ToolEventStage;
+
+use super::CommandTranscript;
+use super::UnifiedExecContext;
+use super::session::UnifiedExecSession;
+
+/// Spawn a background task that continuously reads from the PTY, appends to the
+/// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8
+/// boundaries.
+pub(crate) fn start_streaming_output(
+    session: &UnifiedExecSession,
+    context: &UnifiedExecContext,
+    transcript: Arc<Mutex<CommandTranscript>>,
+) {
+    let mut receiver = session.output_receiver();
+    let session_ref = Arc::clone(&context.session);
+    let turn_ref = Arc::clone(&context.turn);
+    let call_id = context.call_id.clone();
+    let cancellation_token = session.cancellation_token();
+
+    tokio::spawn(async move {
+        let mut pending: Vec<u8> = Vec::new();
+        loop {
+            tokio::select! {
+                _ = cancellation_token.cancelled() => break,
+                result = receiver.recv() => match result {
+                    Ok(chunk) => {
+                        pending.extend_from_slice(&chunk);
+                        while let Some(prefix) = split_valid_utf8_prefix(&mut pending) {
+                            {
+                                let mut guard = transcript.lock().await;
+                                guard.append(&prefix);
+                            }
+
+                            let event = ExecCommandOutputDeltaEvent {
+                                call_id: call_id.clone(),
+                                stream: ExecOutputStream::Stdout,
+                                chunk: prefix,
+                            };
+                            session_ref
+                                .send_event(turn_ref.as_ref(), EventMsg::ExecCommandOutputDelta(event))
+                                .await;
+                        }
+                    }
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                }
+            };
+        }
+    });
+}
+
+/// Spawn a background watcher that waits for the PTY to exit and then emits a
+/// single ExecCommandEnd event with the aggregated transcript.
+#[allow(clippy::too_many_arguments)]
+pub(crate) fn spawn_exit_watcher(
+    session: Arc<UnifiedExecSession>,
+    session_ref: Arc<Session>,
+    turn_ref: Arc<TurnContext>,
+    call_id: String,
+    command: Vec<String>,
+    cwd: PathBuf,
+    process_id: String,
+    transcript: Arc<Mutex<CommandTranscript>>,
+    started_at: Instant,
+) {
+    let exit_token = session.cancellation_token();
+
+    tokio::spawn(async move {
+        exit_token.cancelled().await;
+
+        let exit_code = session.exit_code().unwrap_or(-1);
+        let duration = Instant::now().saturating_duration_since(started_at);
+        emit_exec_end_for_unified_exec(
+            session_ref,
+            turn_ref,
+            call_id,
+            command,
+            cwd,
+            Some(process_id),
+            transcript,
+            String::new(),
+            exit_code,
+            duration,
+        )
+        .await;
+    });
+}
+
+/// Emit an ExecCommandEnd event for a unified exec session, using the transcript
+/// as the primary source of aggregated_output and falling back to the provided
+/// text when the transcript is empty.
+#[allow(clippy::too_many_arguments)]
+pub(crate) async fn emit_exec_end_for_unified_exec(
+    session_ref: Arc<Session>,
+    turn_ref: Arc<TurnContext>,
+    call_id: String,
+    command: Vec<String>,
+    cwd: PathBuf,
+    process_id: Option<String>,
+    transcript: Arc<Mutex<CommandTranscript>>,
+    fallback_output: String,
+    exit_code: i32,
+    duration: Duration,
+) {
+    let aggregated_output = resolve_aggregated_output(&transcript, fallback_output).await;
+    let output = ExecToolCallOutput {
+        exit_code,
+        stdout: StreamOutput::new(aggregated_output.clone()),
+        stderr: StreamOutput::new(String::new()),
+        aggregated_output: StreamOutput::new(aggregated_output),
+        duration,
+        timed_out: false,
+    };
+    let event_ctx = ToolEventCtx::new(session_ref.as_ref(), turn_ref.as_ref(), &call_id, None);
+    let emitter = ToolEmitter::unified_exec(
+        &command,
+        cwd,
+        ExecCommandSource::UnifiedExecStartup,
+        process_id,
+    );
+    emitter
+        .emit(event_ctx, ToolEventStage::Success(output))
+        .await;
+}
+
+fn split_valid_utf8_prefix(buffer: &mut Vec<u8>) -> Option<Vec<u8>> {
+    if buffer.is_empty() {
+        return None;
+    }
+
+    let len = buffer.len();
+    let mut split = len;
+    while split > 0 {
+        if std::str::from_utf8(&buffer[..split]).is_ok() {
+            let prefix = buffer[..split].to_vec();
+            buffer.drain(..split);
+            return Some(prefix);
+        }
+
+        if len - split > 4 {
+            break;
+        }
+        split -= 1;
+    }
+
+    // If no valid UTF-8 prefix was found, emit the first byte so the stream
+    // keeps making progress and the transcript reflects all bytes.
+    let byte = buffer.drain(..1).collect();
+    Some(byte)
+}
+
+async fn resolve_aggregated_output(
+    transcript: &Arc<Mutex<CommandTranscript>>,
+    fallback: String,
+) -> String {
+    let guard = transcript.lock().await;
+    if guard.data.is_empty() {
+        return fallback;
+    }
+
+    String::from_utf8_lossy(&guard.data).to_string()
+}

codex-rs/core/src/unified_exec/mod.rs

@@ -34,6 +34,7 @@ use tokio::sync::Mutex;
 use crate::codex::Session;
 use crate::codex::TurnContext;
 
+mod async_watcher;
 mod errors;
 mod session;
 mod session_manager;
@@ -51,6 +52,24 @@ pub(crate) const MAX_UNIFIED_EXEC_SESSIONS: usize = 64;
 // Send a warning message to the models when it reaches this number of sessions.
 pub(crate) const WARNING_UNIFIED_EXEC_SESSIONS: usize = 60;
 
+#[derive(Debug, Default)]
+pub(crate) struct CommandTranscript {
+    pub data: Vec<u8>,
+}
+
+impl CommandTranscript {
+    pub fn append(&mut self, bytes: &[u8]) {
+        self.data.extend_from_slice(bytes);
+        if self.data.len() > UNIFIED_EXEC_OUTPUT_MAX_BYTES {
+            let excess = self
+                .data
+                .len()
+                .saturating_sub(UNIFIED_EXEC_OUTPUT_MAX_BYTES);
+            self.data.drain(..excess);
+        }
+    }
+}
+
 pub(crate) struct UnifiedExecContext {
     pub session: Arc<Session>,
     pub turn: Arc<TurnContext>,
@@ -92,18 +111,14 @@ pub(crate) struct UnifiedExecResponse {
     pub chunk_id: String,
     pub wall_time: Duration,
     pub output: String,
+    /// Raw bytes returned for this unified exec call before any truncation.
+    pub raw_output: Vec<u8>,
     pub process_id: Option<String>,
     pub exit_code: Option<i32>,
     pub original_token_count: Option<usize>,
     pub session_command: Option<Vec<String>>,
 }
 
-#[derive(Default)]
-pub(crate) struct UnifiedExecSessionManager {
-    session_store: Mutex<SessionStore>,
-}
-
-// Required for mutex sharing.
 #[derive(Default)]
 pub(crate) struct SessionStore {
     sessions: HashMap<String, SessionEntry>,
@@ -115,22 +130,27 @@ impl SessionStore {
         self.reserved_sessions_id.remove(session_id);
         self.sessions.remove(session_id)
     }
+}
 
-    pub(crate) fn clear(&mut self) {
-        self.reserved_sessions_id.clear();
-        self.sessions.clear();
+pub(crate) struct UnifiedExecSessionManager {
+    session_store: Mutex<SessionStore>,
+}
+
+impl Default for UnifiedExecSessionManager {
+    fn default() -> Self {
+        Self {
+            session_store: Mutex::new(SessionStore::default()),
+        }
     }
 }
 
 struct SessionEntry {
-    session: UnifiedExecSession,
+    session: Arc<UnifiedExecSession>,
     session_ref: Arc<Session>,
     turn_ref: Arc<TurnContext>,
     call_id: String,
     process_id: String,
     command: Vec<String>,
-    cwd: PathBuf,
-    started_at: tokio::time::Instant,
     last_used: tokio::time::Instant,
 }
 

codex-rs/core/src/unified_exec/session.rs

@@ -98,19 +98,22 @@ impl UnifiedExecSession {
         let cancellation_token_clone = cancellation_token.clone();
         let output_task = tokio::spawn(async move {
             loop {
-                match receiver.recv().await {
-                    Ok(chunk) => {
-                        let mut guard = buffer_clone.lock().await;
-                        guard.push_chunk(chunk);
-                        drop(guard);
-                        notify_clone.notify_waiters();
+                tokio::select! {
+                    _ = cancellation_token_clone.cancelled() => break,
+                    result = receiver.recv() => match result {
+                        Ok(chunk) => {
+                            let mut guard = buffer_clone.lock().await;
+                            guard.push_chunk(chunk);
+                            drop(guard);
+                            notify_clone.notify_waiters();
+                        }
+                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                        Err(tokio::sync::broadcast::error::RecvError::Closed) => {
+                            cancellation_token_clone.cancel();
+                            break;
+                        }
                     }
-                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
-                    Err(tokio::sync::broadcast::error::RecvError::Closed) => {
-                        cancellation_token_clone.cancel();
-                        break;
-                    }
-                }
+                };
             }
         });
 
@@ -136,6 +139,14 @@ impl UnifiedExecSession {
         }
     }
 
+    pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
+        self.session.output_receiver()
+    }
+
+    pub(super) fn cancellation_token(&self) -> CancellationToken {
+        self.cancellation_token.clone()
+    }
+
     pub(super) fn has_exited(&self) -> bool {
         self.session.has_exited()
     }
@@ -144,6 +155,12 @@ impl UnifiedExecSession {
         self.session.exit_code()
     }
 
+    pub(super) fn terminate(&self) {
+        self.session.terminate();
+        self.cancellation_token.cancel();
+        self.output_task.abort();
+    }
+
     async fn snapshot_output(&self) -> Vec<Vec<u8>> {
         let guard = self.output_buffer.lock().await;
         guard.snapshot()
@@ -246,6 +263,6 @@ impl UnifiedExecSession {
 
 impl Drop for UnifiedExecSession {
     fn drop(&mut self) {
-        self.output_task.abort();
+        self.terminate();
     }
 }

codex-rs/core/src/unified_exec/session_manager.rs

@@ -13,18 +13,12 @@ use tokio_util::sync::CancellationToken;
 use crate::bash::extract_bash_command;
 use crate::codex::Session;
 use crate::codex::TurnContext;
-use crate::exec::ExecToolCallOutput;
-use crate::exec::StreamOutput;
 use crate::exec_env::create_env;
 use crate::exec_policy::create_exec_approval_requirement_for_command;
 use crate::protocol::BackgroundEventEvent;
 use crate::protocol::EventMsg;
-use crate::protocol::ExecCommandSource;
 use crate::sandboxing::ExecEnv;
 use crate::sandboxing::SandboxPermissions;
-use crate::tools::events::ToolEmitter;
-use crate::tools::events::ToolEventCtx;
-use crate::tools::events::ToolEventStage;
 use crate::tools::orchestrator::ToolOrchestrator;
 use crate::tools::runtimes::unified_exec::UnifiedExecRequest as UnifiedExecToolRequest;
 use crate::tools::runtimes::unified_exec::UnifiedExecRuntime;
@@ -33,6 +27,7 @@ use crate::truncate::TruncationPolicy;
 use crate::truncate::approx_token_count;
 use crate::truncate::formatted_truncate_text;
 
+use super::CommandTranscript;
 use super::ExecCommandRequest;
 use super::MAX_UNIFIED_EXEC_SESSIONS;
 use super::SessionEntry;
@@ -43,6 +38,9 @@ use super::UnifiedExecResponse;
 use super::UnifiedExecSessionManager;
 use super::WARNING_UNIFIED_EXEC_SESSIONS;
 use super::WriteStdinRequest;
+use super::async_watcher::emit_exec_end_for_unified_exec;
+use super::async_watcher::spawn_exit_watcher;
+use super::async_watcher::start_streaming_output;
 use super::clamp_yield_time;
 use super::generate_chunk_id;
 use super::resolve_max_tokens;
@@ -135,17 +133,23 @@ impl UnifiedExecSessionManager {
             .await;
 
         let session = match session {
-            Ok(session) => session,
+            Ok(session) => Arc::new(session),
             Err(err) => {
                 self.release_process_id(&request.process_id).await;
                 return Err(err);
             }
         };
 
+        let transcript = Arc::new(tokio::sync::Mutex::new(CommandTranscript::default()));
+        start_streaming_output(&session, context, Arc::clone(&transcript));
+
         let max_tokens = resolve_max_tokens(request.max_output_tokens);
         let yield_time_ms = clamp_yield_time(request.yield_time_ms);
 
         let start = Instant::now();
+        // For the initial exec_command call, we both stream output to events
+        // (via start_streaming_output above) and collect a snapshot here for
+        // the tool response body.
         let OutputHandles {
             output_buffer,
             output_notify,
@@ -163,36 +167,44 @@ impl UnifiedExecSessionManager {
 
         let text = String::from_utf8_lossy(&collected).to_string();
         let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens));
-        let has_exited = session.has_exited();
         let exit_code = session.exit_code();
+        let has_exited = session.has_exited() || exit_code.is_some();
         let chunk_id = generate_chunk_id();
         let process_id = request.process_id.clone();
         if has_exited {
+            // Short‑lived command: emit ExecCommandEnd immediately using the
+            // same helper as the background watcher, so all end events share
+            // one implementation.
             self.release_process_id(&request.process_id).await;
             let exit = exit_code.unwrap_or(-1);
-            Self::emit_exec_end_from_context(
-                context,
-                &request.command,
+            emit_exec_end_for_unified_exec(
+                Arc::clone(&context.session),
+                Arc::clone(&context.turn),
+                context.call_id.clone(),
+                request.command.clone(),
                 cwd,
+                Some(process_id),
+                Arc::clone(&transcript),
                 output.clone(),
                 exit,
                 wall_time,
-                // We always emit the process ID in order to keep consistency between the Begin
-                // event and the End event.
-                Some(process_id),
             )
             .await;
 
             session.check_for_sandbox_denial_with_text(&text).await?;
         } else {
-            // Only store session if not exited.
+            // Long‑lived command: persist the session so write_stdin can reuse
+            // it, and register a background watcher that will emit
+            // ExecCommandEnd when the PTY eventually exits (even if no further
+            // tool calls are made).
             self.store_session(
-                session,
+                Arc::clone(&session),
                 context,
                 &request.command,
                 cwd.clone(),
                 start,
                 process_id,
+                Arc::clone(&transcript),
             )
             .await;
 
@@ -205,6 +217,7 @@ impl UnifiedExecSessionManager {
             chunk_id,
             wall_time,
             output,
+            raw_output: collected,
             process_id: if has_exited {
                 None
             } else {
@@ -238,6 +251,8 @@ impl UnifiedExecSessionManager {
 
         if !request.input.is_empty() {
             Self::send_input(&writer_tx, request.input.as_bytes()).await?;
+            // Give the remote process a brief window to react so that we are
+            // more likely to capture its output in the poll below.
             tokio::time::sleep(Duration::from_millis(100)).await;
         }
 
@@ -259,16 +274,20 @@ impl UnifiedExecSessionManager {
         let original_token_count = approx_token_count(&text);
         let chunk_id = generate_chunk_id();
 
+        // After polling, refresh_session_state tells us whether the PTY is
+        // still alive or has exited and been removed from the store; we thread
+        // that through so the handler can tag TerminalInteraction with an
+        // appropriate process_id and exit_code.
         let status = self.refresh_session_state(process_id.as_str()).await;
-        let (process_id, exit_code, completion_entry, event_call_id) = match status {
+        let (process_id, exit_code, event_call_id) = match status {
             SessionStatus::Alive {
                 exit_code,
                 call_id,
                 process_id,
-            } => (Some(process_id), exit_code, None, call_id),
+            } => (Some(process_id), exit_code, call_id),
             SessionStatus::Exited { exit_code, entry } => {
                 let call_id = entry.call_id.clone();
-                (None, exit_code, Some(*entry), call_id)
+                (None, exit_code, call_id)
             }
             SessionStatus::Unknown => {
                 return Err(UnifiedExecError::UnknownSessionId {
@@ -282,6 +301,7 @@ impl UnifiedExecSessionManager {
             chunk_id,
             wall_time,
             output,
+            raw_output: collected,
             process_id,
             exit_code,
             original_token_count: Some(original_token_count),
@@ -292,12 +312,6 @@ impl UnifiedExecSessionManager {
             Self::emit_waiting_status(&session_ref, &turn_ref, &session_command).await;
         }
 
-        if let (Some(exit), Some(entry)) = (response.exit_code, completion_entry) {
-            let total_duration = Instant::now().saturating_duration_since(entry.started_at);
-            Self::emit_exec_end_from_entry(entry, response.output.clone(), exit, total_duration)
-                .await;
-        }
-
         Ok(response)
     }
 
@@ -371,28 +385,27 @@ impl UnifiedExecSessionManager {
     #[allow(clippy::too_many_arguments)]
     async fn store_session(
         &self,
-        session: UnifiedExecSession,
+        session: Arc<UnifiedExecSession>,
         context: &UnifiedExecContext,
         command: &[String],
         cwd: PathBuf,
         started_at: Instant,
         process_id: String,
+        transcript: Arc<tokio::sync::Mutex<CommandTranscript>>,
     ) {
         let entry = SessionEntry {
-            session,
+            session: Arc::clone(&session),
             session_ref: Arc::clone(&context.session),
             turn_ref: Arc::clone(&context.turn),
             call_id: context.call_id.clone(),
             process_id: process_id.clone(),
             command: command.to_vec(),
-            cwd,
-            started_at,
             last_used: started_at,
         };
         let number_sessions = {
             let mut store = self.session_store.lock().await;
             Self::prune_sessions_if_needed(&mut store);
-            store.sessions.insert(process_id, entry);
+            store.sessions.insert(process_id.clone(), entry);
             store.sessions.len()
         };
 
@@ -405,73 +418,18 @@ impl UnifiedExecSessionManager {
                 )
                 .await;
         };
-    }
 
-    async fn emit_exec_end_from_entry(
-        entry: SessionEntry,
-        aggregated_output: String,
-        exit_code: i32,
-        duration: Duration,
-    ) {
-        let output = ExecToolCallOutput {
-            exit_code,
-            stdout: StreamOutput::new(aggregated_output.clone()),
-            stderr: StreamOutput::new(String::new()),
-            aggregated_output: StreamOutput::new(aggregated_output),
-            duration,
-            timed_out: false,
-        };
-        let event_ctx = ToolEventCtx::new(
-            entry.session_ref.as_ref(),
-            entry.turn_ref.as_ref(),
-            &entry.call_id,
-            None,
-        );
-        let emitter = ToolEmitter::unified_exec(
-            &entry.command,
-            entry.cwd,
-            ExecCommandSource::UnifiedExecStartup,
-            None,
-            Some(entry.process_id.clone()),
-        );
-        emitter
-            .emit(event_ctx, ToolEventStage::Success(output))
-            .await;
-    }
-
-    async fn emit_exec_end_from_context(
-        context: &UnifiedExecContext,
-        command: &[String],
-        cwd: PathBuf,
-        aggregated_output: String,
-        exit_code: i32,
-        duration: Duration,
-        process_id: Option<String>,
-    ) {
-        let output = ExecToolCallOutput {
-            exit_code,
-            stdout: StreamOutput::new(aggregated_output.clone()),
-            stderr: StreamOutput::new(String::new()),
-            aggregated_output: StreamOutput::new(aggregated_output),
-            duration,
-            timed_out: false,
-        };
-        let event_ctx = ToolEventCtx::new(
-            context.session.as_ref(),
-            context.turn.as_ref(),
-            &context.call_id,
-            None,
-        );
-        let emitter = ToolEmitter::unified_exec(
-            command,
+        spawn_exit_watcher(
+            Arc::clone(&session),
+            Arc::clone(&context.session),
+            Arc::clone(&context.turn),
+            context.call_id.clone(),
+            command.to_vec(),
             cwd,
-            ExecCommandSource::UnifiedExecStartup,
-            None,
             process_id,
+            transcript,
+            started_at,
         );
-        emitter
-            .emit(event_ctx, ToolEventStage::Success(output))
-            .await;
     }
 
     async fn emit_waiting_status(
@@ -567,7 +525,7 @@ impl UnifiedExecSessionManager {
         cancellation_token: &CancellationToken,
         deadline: Instant,
     ) -> Vec<u8> {
-        const POST_EXIT_OUTPUT_GRACE: Duration = Duration::from_millis(25);
+        const POST_EXIT_OUTPUT_GRACE: Duration = Duration::from_millis(50);
 
         let mut collected: Vec<u8> = Vec::with_capacity(4096);
         let mut exit_signal_received = cancellation_token.is_cancelled();
@@ -634,7 +592,9 @@ impl UnifiedExecSessionManager {
             .collect();
 
         if let Some(session_id) = Self::session_id_to_prune_from_meta(&meta) {
-            store.remove(&session_id);
+            if let Some(entry) = store.remove(&session_id) {
+                entry.session.terminate();
+            }
             return true;
         }
 
@@ -671,8 +631,17 @@ impl UnifiedExecSessionManager {
     }
 
     pub(crate) async fn terminate_all_sessions(&self) {
-        let mut sessions = self.session_store.lock().await;
-        sessions.clear();
+        let entries: Vec<SessionEntry> = {
+            let mut sessions = self.session_store.lock().await;
+            let entries: Vec<SessionEntry> =
+                sessions.sessions.drain().map(|(_, entry)| entry).collect();
+            sessions.reserved_sessions_id.clear();
+            entries
+        };
+
+        for entry in entries {
+            entry.session.terminate();
+        }
     }
 }