Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions codex-rs/analytics/src/events.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1140,6 +1140,7 @@ pub(crate) struct CodexTurnEventParams {
pub(crate) after_last_sampling_ms: u64,
pub(crate) sampling_request_count: u32,
pub(crate) sampling_retry_count: u32,
pub(crate) tools_change_count: u32,
pub(crate) duration_ms: Option<u64>,
pub(crate) started_at: Option<u64>,
pub(crate) completed_at: Option<u64>,
Expand Down
1 change: 1 addition & 0 deletions codex-rs/analytics/src/facts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ pub struct TurnProfile {
pub after_last_sampling_ms: u64,
pub sampling_request_count: u32,
pub sampling_retry_count: u32,
pub tools_change_count: u32,
}

#[derive(Clone)]
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/analytics/src/reducer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3667,6 +3667,7 @@ fn codex_turn_event_params(
after_last_sampling_ms,
sampling_request_count,
sampling_retry_count,
tools_change_count,
} = profile;
let token_usage = turn_state.token_usage.clone();
let codex_error = turn_state.codex_error.as_ref();
Expand Down Expand Up @@ -3750,6 +3751,7 @@ fn codex_turn_event_params(
after_last_sampling_ms,
sampling_request_count,
sampling_retry_count,
tools_change_count,
duration_ms: completed.duration_ms,
started_at,
completed_at: Some(completed.completed_at),
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/analytics/src/tests/suite/turns/events_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ fn turn_event_serializes_expected_shape() {
after_last_sampling_ms: 94,
sampling_request_count: 2,
sampling_retry_count: 1,
tools_change_count: 2,
duration_ms: Some(1234),
started_at: Some(455),
completed_at: Some(456),
Expand Down Expand Up @@ -264,6 +265,7 @@ fn turn_event_serializes_expected_shape() {
"after_last_sampling_ms": 94,
"sampling_request_count": 2,
"sampling_retry_count": 1,
"tools_change_count": 2,
"duration_ms": 1234,
"started_at": 455,
"completed_at": 456
Expand Down
1 change: 1 addition & 0 deletions codex-rs/analytics/src/tests/support.rs
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,7 @@ pub(super) fn sample_turn_profile() -> TurnProfile {
after_last_sampling_ms: 94,
sampling_request_count: 2,
sampling_retry_count: 1,
tools_change_count: 2,
}
}

Expand Down
15 changes: 15 additions & 0 deletions codex-rs/core/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout_trace::InferenceTraceAttempt;
use codex_rollout_trace::InferenceTraceContext;
use codex_tools::ToolSpec;
use codex_tools::create_tools_json_for_responses_lite;
use codex_tools::create_tools_raw_json_for_responses_api;
use eventsource_stream::Event;
Expand Down Expand Up @@ -217,6 +218,8 @@ struct ModelClientState {
disable_websockets: AtomicBool,
agent_identity_session_fallback: AgentIdentitySessionFallback,
cached_websocket_session: StdMutex<WebsocketSession>,
/// Last full tool list used for sampling, retained across turns and connection resets.
last_inference_tools: StdMutex<Option<Arc<[ToolSpec]>>>,
}

enum ClientRouting {
Expand Down Expand Up @@ -533,6 +536,7 @@ impl ModelClient {
disable_websockets: AtomicBool::new(false),
agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
last_inference_tools: StdMutex::new(None),
}),
agent_identity_policy,
prompt_cache_key_override: None,
Expand Down Expand Up @@ -1350,6 +1354,17 @@ impl Drop for ModelClientSession {
}

impl ModelClientSession {
pub(crate) fn inference_tools_changed(&self, tools: &Arc<[ToolSpec]>) -> bool {
let previous = self
.client
.state
.last_inference_tools
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.replace(Arc::clone(tools));
previous.is_some_and(|previous| previous != *tools)
}

#[allow(clippy::too_many_arguments)]
/// Builds shared Responses API transport options and request-body options.
///
Expand Down
55 changes: 55 additions & 0 deletions codex-rs/core/src/client_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ use codex_rollout_trace::RawTraceEventPayload;
use codex_rollout_trace::RolloutTrace;
use codex_rollout_trace::TraceWriter;
use codex_rollout_trace::replay_bundle;
use codex_tools::JsonSchema;
use codex_tools::ResponsesApiTool;
use codex_tools::ToolSpec;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use serde_json::json;
Expand Down Expand Up @@ -2130,3 +2133,55 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow:
assert_eq!(recorded["output_items"], serde_json::to_value(&delivered)?);
Ok(())
}

#[tokio::test]
async fn inference_tools_changes_follow_full_specs_across_turns() {
let client = test_model_client(SessionSource::Cli);
let telemetry = test_session_telemetry();
let model_info = test_model_info();
let alpha = ResponsesApiTool {
name: "alpha".into(),
description: "Original".into(),
strict: false,
defer_loading: None,
parameters: JsonSchema::default(),
output_schema: None,
};
let beta = ResponsesApiTool {
name: "beta".into(),
..alpha.clone()
};
let mut output_changed = alpha.clone();
output_changed.output_schema = Some(json!({"type": "object"}).into());
let mut params_changed = beta.clone();
params_changed.parameters = JsonSchema::string(Some("Changed parameters".into()));
let mut session = client.new_session();
for (index, (tools, expected)) in [
(vec![alpha.clone()], false),
(vec![output_changed.clone()], true),
(vec![output_changed.clone(), beta.clone()], true),
(vec![output_changed, beta.clone()], false),
(vec![beta, alpha.clone()], true),
(vec![params_changed, alpha], true),
(vec![], true),
(vec![], false),
]
.into_iter()
.enumerate()
{
if index == 4 {
session.try_switch_fallback_transport(&telemetry, &model_info);
drop(session);
session = client.new_session();
}
let specs = tools
.into_iter()
.map(ToolSpec::Function)
.collect::<Arc<[_]>>();
assert_eq!(
session.inference_tools_changed(&specs),
expected,
"step {index}"
);
}
}
3 changes: 3 additions & 0 deletions codex-rs/core/src/session/turn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2568,6 +2568,9 @@ async fn try_run_sampling_request(
turn_context.provider.info().name.as_str(),
);
let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling();
if client_session.inference_tools_changed(&step_context.tool_router.model_visible_specs()) {
turn_context.turn_timing_state.record_tools_change();
}
// Do not enter this span: overlapping tools must not retain it past sampling.
let sampling_span = trace_span!(
"codex.sampling",
Expand Down
9 changes: 9 additions & 0 deletions codex-rs/core/src/turn_timing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ struct TurnProfileState {
pending_idle_after_sampling: Duration,
sampling_request_count: u32,
sampling_retry_count: u32,
tools_change_count: u32,
completed_profile: Option<TurnProfile>,
}

Expand Down Expand Up @@ -155,6 +156,13 @@ impl TurnTimingState {
self.profile_state().record_sampling_retry();
}

pub(crate) fn record_tools_change(&self) {
let mut profile = self.profile_state();
if profile.completed_profile.is_none() && profile.started_at.is_some() {
profile.tools_change_count = profile.tools_change_count.saturating_add(1);
}
}

pub(crate) fn begin_compaction(self: &Arc<Self>) -> TurnProfileTimingGuard {
let active = self.profile_state().begin_compaction(Instant::now());
TurnProfileTimingGuard {
Expand Down Expand Up @@ -324,6 +332,7 @@ impl TurnProfileState {
after_last_sampling_ms: duration_to_u64_ms(after_last_sampling),
sampling_request_count: self.sampling_request_count,
sampling_retry_count: self.sampling_retry_count,
tools_change_count: self.tools_change_count,
};
let total_ms = self
.started_at
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/turn_timing_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,7 @@ fn turn_profile_breaks_down_sampling_blocking_and_retry_overhead() {
after_last_sampling_ms: 100,
sampling_request_count: 2,
sampling_retry_count: 1,
tools_change_count: 0,
}
);
}
Expand Down Expand Up @@ -270,6 +271,7 @@ fn turn_profile_counts_compaction_as_an_exclusive_phase() {
after_last_sampling_ms: 100,
sampling_request_count: 1,
sampling_retry_count: 0,
tools_change_count: 0,
}
);
}
Expand Down
Loading