agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rollout_trace] Trace tool and code-mode boundaries (#18878)

Browse Source 
## Summary

Extends rollout tracing across tool dispatch and code-mode runtime
boundaries. This records canonical tool-call lifecycle events and links
code-mode execution/wait operations back to the model-visible calls that
caused them.

## Stack

This is PR 3/5 in the rollout trace stack.

- [#18876](https://github.com/openai/codex/pull/18876): Add rollout
trace crate
- [#18877](https://github.com/openai/codex/pull/18877): Record core
session rollout traces
- [#18878](https://github.com/openai/codex/pull/18878): Trace tool and
code-mode boundaries
- [#18879](https://github.com/openai/codex/pull/18879): Trace sessions
and multi-agent edges
- [#18880](https://github.com/openai/codex/pull/18880): Add debug trace
reduction command

## Review Notes

This PR is about attribution. Reviewers should focus on whether direct
tool calls, code-mode-originated tool calls, waits, outputs, and
cancellation boundaries are recorded with enough source information for
deterministic reduction without coupling the reducer to live runtime
internals.

The stack remains valid after this layer: tool and code-mode traces
reduce through the existing crate model, while the broader session and
multi-agent relationships are added in the next PR.
This commit is contained in:
cassirer-openai
2026-04-23 12:22:11 -07:00
committed by GitHub
parent ff22982d75
commit 6d09b6752d
30 changed files with 1553 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/code_mode/execute_handler.rs
+22 -3
View File
@@ -9,6 +9,7 @@ use super::ExecContext;
use super::PUBLIC_TOOL_NAME;
use super::build_enabled_tools;
use super::handle_runtime_response;
use super::is_exec_tool_name;
pub struct CodeModeExecuteHandler;
@@ -30,12 +31,23 @@ impl CodeModeExecuteHandler {
.code_mode_service
.stored_values()
.await;
// Allocate before starting V8 so the trace can create the parent
// CodeCell before model-authored JavaScript issues nested tool calls.
let runtime_cell_id = exec.session.services.code_mode_service.allocate_cell_id();
let code_cell_trace = exec.session.services.rollout_trace.start_code_cell_trace(
exec.session.conversation_id,
exec.turn.sub_id.as_str(),
runtime_cell_id.as_str(),
call_id.as_str(),
args.code.as_str(),
);
let started_at = std::time::Instant::now();
let response = exec
.session
.services
.code_mode_service
.execute(codex_code_mode::ExecuteRequest {
cell_id: runtime_cell_id,
tool_call_id: call_id,
enabled_tools,
source: args.code,
@@ -45,6 +57,15 @@ impl CodeModeExecuteHandler {
})
.await
.map_err(FunctionCallError::RespondToModel)?;
// Record the raw runtime boundary. The model-visible custom-tool output
// is produced by `handle_runtime_response` and later linked through
// `CodeCell.output_item_ids` in the reduced trace.
code_cell_trace.record_initial_response(&response);
// Yielded cells keep running, so terminal lifecycle is only emitted
// here when the first response also ended the runtime.
if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) {
code_cell_trace.record_ended(&response);
}
handle_runtime_response(&exec, response, args.max_output_tokens, started_at)
.await
.map_err(FunctionCallError::RespondToModel)
@@ -73,9 +94,7 @@ impl ToolHandler for CodeModeExecuteHandler {
} = invocation;
match payload {
ToolPayload::Custom { input }
if tool_name.namespace.is_none() && tool_name.name.as_str() == PUBLIC_TOOL_NAME =>
{
ToolPayload::Custom { input } if is_exec_tool_name(&tool_name) => {
self.execute(session, turn, call_id, input).await
}
_ => Err(FunctionCallError::RespondToModel(format!(
codex-rs/core/src/tools/code_mode/mod.rs
+29 -9
View File
@@ -7,6 +7,7 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeTurnHost;
use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
@@ -45,6 +46,11 @@ pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME;
pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode::WAIT_TOOL_NAME;
pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = codex_code_mode::DEFAULT_WAIT_YIELD_TIME_MS;
/// Returns true for the un-namespaced code-mode `exec` tool.
pub(crate) fn is_exec_tool_name(tool_name: &ToolName) -> bool {
tool_name.namespace.is_none() && tool_name.name == PUBLIC_TOOL_NAME
}
#[derive(Clone)]
pub(crate) struct ExecContext {
pub(super) session: Arc<Session>,
@@ -73,6 +79,10 @@ impl CodeModeService {
self.inner.replace_stored_values(values).await;
}
pub(crate) fn allocate_cell_id(&self) -> String {
self.inner.allocate_cell_id()
}
pub(crate) async fn execute(
&self,
request: codex_code_mode::ExecuteRequest,
@@ -83,7 +93,7 @@ impl CodeModeService {
pub(crate) async fn wait(
&self,
request: codex_code_mode::WaitRequest,
) -> Result<RuntimeResponse, String> {
) -> Result<codex_code_mode::WaitOutcome, String> {
self.inner.wait(request).await
}
@@ -118,15 +128,13 @@ struct CoreTurnHost {
impl CodeModeTurnHost for CoreTurnHost {
async fn invoke_tool(
&self,
tool_name: ToolName,
input: Option<JsonValue>,
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
call_nested_tool(
self.exec.clone(),
self.tool_runtime.clone(),
tool_name,
input,
invocation,
cancellation_token,
)
.await
@@ -302,11 +310,16 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter {
async fn call_nested_tool(
exec: ExecContext,
tool_runtime: ToolCallRuntime,
tool_name: ToolName,
input: Option<JsonValue>,
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> Result<JsonValue, FunctionCallError> {
if tool_name.namespace.is_none() && tool_name.name == PUBLIC_TOOL_NAME {
let CodeModeNestedToolCall {
cell_id,
runtime_tool_call_id,
tool_name,
input,
} = invocation;
if is_exec_tool_name(&tool_name) {
return Err(FunctionCallError::RespondToModel(format!(
"{PUBLIC_TOOL_NAME} cannot invoke itself"
)));
@@ -339,7 +352,14 @@ async fn call_nested_tool(
payload,
};
let result = tool_runtime
.handle_tool_call_with_source(call, ToolCallSource::CodeMode, cancellation_token)
.handle_tool_call_with_source(
call,
ToolCallSource::CodeMode {
cell_id,
runtime_tool_call_id,
},
cancellation_token,
)
.await?;
Ok(result.code_mode_result())
}
codex-rs/core/src/tools/code_mode/wait_handler.rs
+23 -2
View File
@@ -61,7 +61,7 @@ impl ToolHandler for CodeModeWaitHandler {
let args: ExecWaitArgs = parse_arguments(&arguments)?;
let exec = ExecContext { session, turn };
let started_at = std::time::Instant::now();
let response = exec
let wait_response = exec
.session
.services
.code_mode_service
@@ -72,7 +72,28 @@ impl ToolHandler for CodeModeWaitHandler {
})
.await
.map_err(FunctionCallError::RespondToModel)?;
handle_runtime_response(&exec, response, args.max_tokens, started_at)
if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response
&& !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. })
{
// Only a live-cell wait can close a CodeCell. A missing
// cell is still an ordinary `wait` tool result, but there
// is no runtime object for the reducer to complete.
let runtime_cell_id = match response {
codex_code_mode::RuntimeResponse::Yielded { cell_id, .. }
| codex_code_mode::RuntimeResponse::Terminated { cell_id, .. }
| codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id,
};
exec.session
.services
.rollout_trace
.code_cell_trace_context(
exec.session.conversation_id,
exec.turn.sub_id.as_str(),
runtime_cell_id,
)
.record_ended(response);
}
handle_runtime_response(&exec, wait_response.into(), args.max_tokens, started_at)
.await
.map_err(FunctionCallError::RespondToModel)
}