From e2eb2bbbb721474f814afee92bf813110e33edcd Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 10 Jul 2026 00:06:25 +0000 Subject: [PATCH] codex: address PR review feedback (#31707) --- .../src/client/reqwest_http_client.rs | 133 ++++++++++-------- codex-rs/exec-server/src/local_process.rs | 124 ++++++++-------- codex-rs/exec-server/src/server/handler.rs | 33 ++--- 3 files changed, 153 insertions(+), 137 deletions(-) diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index f8033f4046..3f10d3782f 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -19,6 +19,7 @@ use reqwest::Url; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use reqwest::header::HeaderValue; +use tokio_util::sync::CancellationToken; use tracing::Instrument; use super::HttpResponseBodyStream; @@ -209,67 +210,89 @@ impl ReqwestHttpRequestRunner { )) } + #[tracing::instrument( + name = "codex.exec_server.http_response_body", + parent = None, + skip(pending_stream, notifications, shutdown, initiating_span), + follows_from = [&initiating_span], + fields( + otel.kind = "internal", + exec_server.http_request_id = %pending_stream.request_id, + result = tracing::field::Empty, + error.type = tracing::field::Empty, + ) + )] pub(crate) async fn stream_body( pending_stream: PendingReqwestHttpBodyStream, notifications: RpcNotificationSender, + shutdown: CancellationToken, + initiating_span: tracing::Span, ) { - let PendingReqwestHttpBodyStream { - request_id, - response, - } = pending_stream; - let mut seq = 1; - let mut body = response.bytes_stream(); - while let Some(chunk) = body.next().await { - match chunk { - Ok(bytes) => { - if !send_body_delta( - ¬ifications, - HttpRequestBodyDeltaNotification { - request_id: request_id.clone(), - seq, - delta: bytes.to_vec().into(), - done: false, - error: None, - }, - ) - .await - { - tracing::Span::current().record("result", "disconnected"); - return; - } - seq += 1; - } - Err(error) => { - tracing::Span::current().record("result", "error"); - tracing::Span::current().record("error.type", "response_body"); - let _ = send_body_delta( - ¬ifications, - HttpRequestBodyDeltaNotification { - request_id, - seq, - delta: Vec::new().into(), - done: true, - error: Some(error.to_string()), - }, - ) - .await; - return; - } + tokio::select! { + _ = shutdown.cancelled() => { + tracing::Span::current().record("result", "cancelled"); } - } + _ = async move { + let PendingReqwestHttpBodyStream { + request_id, + response, + } = pending_stream; + let mut seq = 1; + let mut body = response.bytes_stream(); + while let Some(chunk) = body.next().await { + match chunk { + Ok(bytes) => { + if !send_body_delta( + ¬ifications, + HttpRequestBodyDeltaNotification { + request_id: request_id.clone(), + seq, + delta: bytes.to_vec().into(), + done: false, + error: None, + }, + ) + .await + { + tracing::Span::current().record("result", "disconnected"); + return; + } + seq += 1; + } + Err(error) => { + tracing::Span::current().record("result", "error"); + tracing::Span::current().record("error.type", "response_body"); + let _ = send_body_delta( + ¬ifications, + HttpRequestBodyDeltaNotification { + request_id, + seq, + delta: Vec::new().into(), + done: true, + error: Some(error.to_string()), + }, + ) + .await; + return; + } + } + } - let sent = send_body_delta( - ¬ifications, - HttpRequestBodyDeltaNotification { - request_id, - seq, - delta: Vec::new().into(), - done: true, - error: None, - }, - ) - .await; - tracing::Span::current().record("result", if sent { "success" } else { "disconnected" }); + let sent = send_body_delta( + ¬ifications, + HttpRequestBodyDeltaNotification { + request_id, + seq, + delta: Vec::new().into(), + done: true, + error: None, + }, + ) + .await; + tracing::Span::current() + .record("result", if sent { "success" } else { "disconnected" }); + } => {} + } } fn build_headers(headers: Vec) -> Result { diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 3199358141..945dd01432 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -22,7 +22,6 @@ use tokio::sync::Mutex; use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::sync::watch; -use tracing::Instrument; use crate::ExecBackend; use crate::ExecBackendFuture; @@ -342,69 +341,39 @@ impl LocalProcess { })), ); } - let (stdout_stream, stdout_stream_name) = if params.tty { - (ExecOutputStream::Pty, "pty") + let stdout_stream = if params.tty { + ExecOutputStream::Pty } else { - (ExecOutputStream::Stdout, "stdout") + ExecOutputStream::Stdout }; - let stdout_span = tracing::info_span!( - parent: None, - "codex.exec_server.process_output", - otel.kind = "internal", - exec_server.process_id = %process_id, - exec_server.output_stream = stdout_stream_name, - ); - stdout_span.follows_from(tracing::Span::current()); - tokio::spawn( - stream_output( - process_id.clone(), - stdout_stream, - spawned.stdout_rx, - Arc::clone(&self.inner), - Arc::clone(&output_notify), - ) - .instrument(stdout_span), - ); - let (stderr_stream, stderr_stream_name) = if params.tty { - (ExecOutputStream::Pty, "pty") + tokio::spawn(stream_output( + process_id.clone(), + stdout_stream, + spawned.stdout_rx, + Arc::clone(&self.inner), + Arc::clone(&output_notify), + tracing::Span::current(), + )); + let stderr_stream = if params.tty { + ExecOutputStream::Pty } else { - (ExecOutputStream::Stderr, "stderr") + ExecOutputStream::Stderr }; - let stderr_span = tracing::info_span!( - parent: None, - "codex.exec_server.process_output", - otel.kind = "internal", - exec_server.process_id = %process_id, - exec_server.output_stream = stderr_stream_name, - ); - stderr_span.follows_from(tracing::Span::current()); - tokio::spawn( - stream_output( - process_id.clone(), - stderr_stream, - spawned.stderr_rx, - Arc::clone(&self.inner), - Arc::clone(&output_notify), - ) - .instrument(stderr_span), - ); - let wait_span = tracing::info_span!( - parent: None, - "codex.exec_server.process_wait", - otel.kind = "internal", - exec_server.process_id = %process_id, - process.exit_code = tracing::field::Empty, - ); - wait_span.follows_from(tracing::Span::current()); - tokio::spawn( - watch_exit( - process_id.clone(), - spawned.exit_rx, - Arc::clone(&self.inner), - output_notify, - ) - .instrument(wait_span), - ); + tokio::spawn(stream_output( + process_id.clone(), + stderr_stream, + spawned.stderr_rx, + Arc::clone(&self.inner), + Arc::clone(&output_notify), + tracing::Span::current(), + )); + tokio::spawn(watch_exit( + process_id.clone(), + spawned.exit_rx, + Arc::clone(&self.inner), + output_notify, + tracing::Span::current(), + )); Ok((ExecResponse { process_id }, wake_tx, events)) } @@ -794,13 +763,33 @@ fn map_handler_error(error: JSONRPCErrorError) -> ExecServerError { } } +#[tracing::instrument( + name = "codex.exec_server.process_output", + parent = None, + skip(receiver, inner, output_notify, initiating_span), + follows_from = [&initiating_span], + fields( + otel.kind = "internal", + exec_server.process_id = %process_id, + exec_server.output_stream = tracing::field::Empty, + ) +)] async fn stream_output( process_id: ProcessId, stream: ExecOutputStream, mut receiver: tokio::sync::mpsc::Receiver>, inner: Arc, output_notify: Arc, + initiating_span: tracing::Span, ) { + tracing::Span::current().record( + "exec_server.output_stream", + match stream { + ExecOutputStream::Stdout => "stdout", + ExecOutputStream::Stderr => "stderr", + ExecOutputStream::Pty => "pty", + }, + ); while let Some(chunk) = receiver.recv().await { let _chunk_len = chunk.len(); let notification = { @@ -852,11 +841,23 @@ async fn stream_output( finish_output_stream(process_id, inner).await; } +#[tracing::instrument( + name = "codex.exec_server.process_wait", + parent = None, + skip(exit_rx, inner, output_notify, initiating_span), + follows_from = [&initiating_span], + fields( + otel.kind = "internal", + exec_server.process_id = %process_id, + process.exit_code = tracing::field::Empty, + ) +)] async fn watch_exit( process_id: ProcessId, exit_rx: tokio::sync::oneshot::Receiver, inner: Arc, output_notify: Arc, + initiating_span: tracing::Span, ) { let exit_code = exit_rx.await.unwrap_or(-1); tracing::Span::current().record("process.exit_code", exit_code); @@ -1358,6 +1359,7 @@ mod tests { stdout_rx, Arc::clone(&backend.inner), Arc::clone(&output_notify), + tracing::Span::current(), )); tokio::spawn(stream_output( process_id.clone(), @@ -1365,12 +1367,14 @@ mod tests { stderr_rx, Arc::clone(&backend.inner), Arc::clone(&output_notify), + tracing::Span::current(), )); tokio::spawn(watch_exit( process_id.clone(), exit_rx, Arc::clone(&backend.inner), output_notify, + tracing::Span::current(), )); TestProcess { diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index c4a7a37caa..62565349cf 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -10,7 +10,6 @@ use std::collections::HashSet; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; -use tracing::Instrument; use crate::ExecServerRuntimePaths; use crate::client::http_client::PendingReqwestHttpBodyStream; @@ -391,27 +390,17 @@ impl ExecServerHandler { let handler = Arc::clone(self); let notifications = self.notifications.clone(); let shutdown = self.background_task_shutdown.clone(); - let stream_span = tracing::info_span!( - parent: None, - "codex.exec_server.http_response_body", - otel.kind = "internal", - exec_server.http_request_id = request_id, - result = tracing::field::Empty, - error.type = tracing::field::Empty, - ); - stream_span.follows_from(tracing::Span::current()); - self.background_tasks.spawn( - async move { - tokio::select! { - _ = shutdown.cancelled() => { - tracing::Span::current().record("result", "cancelled"); - } - _ = ReqwestHttpRequestRunner::stream_body(pending_stream, notifications) => {} - } - handler.release_http_body_stream(&finished_request_id).await; - } - .instrument(stream_span), - ); + let initiating_span = tracing::Span::current(); + self.background_tasks.spawn(async move { + ReqwestHttpRequestRunner::stream_body( + pending_stream, + notifications, + shutdown, + initiating_span, + ) + .await; + handler.release_http_body_stream(&finished_request_id).await; + }); } async fn release_http_body_stream(&self, request_id: &str) {