codex: address PR review feedback (#31707)

This commit is contained in:
Adam Perry
2026-07-10 00:06:25 +00:00
parent e62c19c428
commit e2eb2bbbb7
3 changed files with 153 additions and 137 deletions

View File

@@ -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(
&notifications,
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(
&notifications,
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(
&notifications,
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(
&notifications,
HttpRequestBodyDeltaNotification {
request_id,
seq,
delta: Vec::new().into(),
done: true,
error: Some(error.to_string()),
},
)
.await;
return;
}
}
}
let sent = send_body_delta(
&notifications,
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(
&notifications,
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<HttpHeader>) -> Result<HeaderMap, JSONRPCErrorError> {

View File

@@ -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<Vec<u8>>,
inner: Arc<Inner>,
output_notify: Arc<Notify>,
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<i32>,
inner: Arc<Inner>,
output_notify: Arc<Notify>,
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 {

View File

@@ -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) {