codex_rollout_trace/
compaction.rs

1//! Hot-path helpers for recording upstream remote compaction attempts.
2//!
3//! Remote compaction is a model-facing request with a different semantic role
4//! from normal sampling. Keeping the no-op capable trace handle in this crate
5//! lets `codex-core` record exact endpoint payloads without owning trace schema
6//! details.
7
8use std::fmt::Display;
9use std::sync::Arc;
10use std::sync::atomic::AtomicU64;
11use std::sync::atomic::Ordering;
12
13use codex_protocol::models::ResponseItem;
14use serde::Serialize;
15use serde_json::Value as JsonValue;
16use tracing::warn;
17
18use crate::inference::trace_response_item_json;
19use crate::model::AgentThreadId;
20use crate::model::CodexTurnId;
21use crate::model::CompactionId;
22use crate::model::CompactionRequestId;
23use crate::payload::RawPayloadKind;
24use crate::raw_event::RawTraceEventContext;
25use crate::raw_event::RawTraceEventPayload;
26use crate::writer::TraceWriter;
27
28static NEXT_COMPACTION_REQUEST: AtomicU64 = AtomicU64::new(1);
29
30/// Turn-local remote compaction tracing context.
31///
32/// A compaction can retry its upstream request before installing one checkpoint. The context
33/// owns the stable checkpoint ID; each request attempt gets a separate request ID.
34#[derive(Clone, Debug)]
35pub struct CompactionTraceContext {
36    state: CompactionTraceContextState,
37}
38
39#[derive(Clone, Debug)]
40enum CompactionTraceContextState {
41    Disabled,
42    Enabled(EnabledCompactionTraceContext),
43}
44
45#[derive(Clone, Debug)]
46struct EnabledCompactionTraceContext {
47    writer: Arc<TraceWriter>,
48    thread_id: AgentThreadId,
49    codex_turn_id: CodexTurnId,
50    compaction_id: CompactionId,
51    model: String,
52    provider_name: String,
53}
54
55/// One upstream request attempt made while computing a compaction checkpoint.
56#[derive(Clone, Debug)]
57pub struct CompactionTraceAttempt {
58    state: CompactionTraceAttemptState,
59}
60
61#[derive(Clone, Debug)]
62enum CompactionTraceAttemptState {
63    Disabled,
64    Enabled(EnabledCompactionTraceAttempt),
65}
66
67#[derive(Clone, Debug)]
68struct EnabledCompactionTraceAttempt {
69    context: EnabledCompactionTraceContext,
70    compaction_request_id: CompactionRequestId,
71}
72
73#[derive(Serialize)]
74struct TracedCompactionCompleted {
75    output_items: Vec<JsonValue>,
76}
77
78/// History replacement checkpoint persisted when compaction installs new live history.
79///
80/// The checkpoint keeps compaction separate from ordinary sampling snapshots:
81/// `input_history` is the live thread history selected for compaction, while
82/// `replacement_history` is what future prompts may carry after the checkpoint.
83#[derive(Serialize)]
84pub struct CompactionCheckpointTracePayload<'a> {
85    pub input_history: &'a [ResponseItem],
86    pub replacement_history: &'a [ResponseItem],
87}
88
89impl CompactionTraceContext {
90    /// Builds a context that accepts trace calls and records nothing.
91    pub fn disabled() -> Self {
92        Self {
93            state: CompactionTraceContextState::Disabled,
94        }
95    }
96
97    /// Builds an enabled context for upstream attempts that compute one checkpoint.
98    pub fn enabled(
99        writer: Arc<TraceWriter>,
100        thread_id: AgentThreadId,
101        codex_turn_id: CodexTurnId,
102        compaction_id: CompactionId,
103        model: String,
104        provider_name: String,
105    ) -> Self {
106        Self {
107            state: CompactionTraceContextState::Enabled(EnabledCompactionTraceContext {
108                writer,
109                thread_id,
110                codex_turn_id,
111                compaction_id,
112                model,
113                provider_name,
114            }),
115        }
116    }
117
118    /// Returns whether this context records compaction traces.
119    pub fn is_enabled(&self) -> bool {
120        matches!(self.state, CompactionTraceContextState::Enabled(_))
121    }
122
123    /// Starts a new upstream attempt and records the exact compact endpoint request.
124    pub fn start_attempt(&self, request: &impl Serialize) -> CompactionTraceAttempt {
125        let CompactionTraceContextState::Enabled(context) = &self.state else {
126            return CompactionTraceAttempt::disabled();
127        };
128
129        let attempt = CompactionTraceAttempt {
130            state: CompactionTraceAttemptState::Enabled(EnabledCompactionTraceAttempt {
131                context: context.clone(),
132                compaction_request_id: next_compaction_request_id(),
133            }),
134        };
135        attempt.record_started(request);
136        attempt
137    }
138
139    /// Records the point where compacted history becomes the live thread history.
140    ///
141    /// The checkpoint belongs to the same semantic compaction lifecycle as the
142    /// compact endpoint attempts, so the context reuses its stable compaction ID.
143    pub fn record_installed(&self, checkpoint: &CompactionCheckpointTracePayload<'_>) {
144        let CompactionTraceContextState::Enabled(context) = &self.state else {
145            return;
146        };
147        let checkpoint_payload = match context
148            .writer
149            .write_json_payload(RawPayloadKind::CompactionCheckpoint, checkpoint)
150        {
151            Ok(payload_ref) => payload_ref,
152            Err(err) => {
153                warn!("failed to write rollout trace payload: {err:#}");
154                return;
155            }
156        };
157
158        let event_context = RawTraceEventContext {
159            thread_id: Some(context.thread_id.clone()),
160            codex_turn_id: Some(context.codex_turn_id.clone()),
161        };
162        if let Err(err) = context.writer.append_with_context(
163            event_context,
164            RawTraceEventPayload::CompactionInstalled {
165                compaction_id: context.compaction_id.clone(),
166                checkpoint_payload,
167            },
168        ) {
169            warn!("failed to append rollout trace event: {err:#}");
170        }
171    }
172}
173
174impl CompactionTraceAttempt {
175    /// Builds an attempt that records nothing.
176    fn disabled() -> Self {
177        Self {
178            state: CompactionTraceAttemptState::Disabled,
179        }
180    }
181
182    fn record_started(&self, request: &impl Serialize) {
183        let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
184            return;
185        };
186        let Some(request_payload) = write_json_payload_best_effort(
187            &attempt.context.writer,
188            RawPayloadKind::CompactionRequest,
189            request,
190        ) else {
191            return;
192        };
193
194        append_with_context_best_effort(
195            &attempt.context,
196            RawTraceEventPayload::CompactionRequestStarted {
197                compaction_id: attempt.context.compaction_id.clone(),
198                compaction_request_id: attempt.compaction_request_id.clone(),
199                thread_id: attempt.context.thread_id.clone(),
200                codex_turn_id: attempt.context.codex_turn_id.clone(),
201                model: attempt.context.model.clone(),
202                provider_name: attempt.context.provider_name.clone(),
203                request_payload,
204            },
205        );
206    }
207
208    /// Records the non-streaming compact endpoint response payload.
209    ///
210    /// Compaction responses use the same response-item preservation rules as
211    /// inference streams: traces are evidence, while normal ResponseItem
212    /// serialization is shaped for future request construction.
213    pub fn record_completed(&self, output_items: &[ResponseItem]) {
214        let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
215            return;
216        };
217        let response_payload = TracedCompactionCompleted {
218            output_items: output_items.iter().map(trace_response_item_json).collect(),
219        };
220        let Some(response_payload) = write_json_payload_best_effort(
221            &attempt.context.writer,
222            RawPayloadKind::CompactionResponse,
223            &response_payload,
224        ) else {
225            return;
226        };
227
228        append_with_context_best_effort(
229            &attempt.context,
230            RawTraceEventPayload::CompactionRequestCompleted {
231                compaction_id: attempt.context.compaction_id.clone(),
232                compaction_request_id: attempt.compaction_request_id.clone(),
233                response_payload,
234            },
235        );
236    }
237
238    /// Records the compact endpoint result without forcing callers to branch on trace events.
239    pub fn record_result<E: Display>(&self, result: Result<&[ResponseItem], E>) {
240        match result {
241            Ok(output_items) => self.record_completed(output_items),
242            Err(err) => self.record_failed(err),
243        }
244    }
245
246    /// Records pre-response failures from the compact endpoint.
247    pub fn record_failed(&self, error: impl Display) {
248        let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
249            return;
250        };
251        append_with_context_best_effort(
252            &attempt.context,
253            RawTraceEventPayload::CompactionRequestFailed {
254                compaction_id: attempt.context.compaction_id.clone(),
255                compaction_request_id: attempt.compaction_request_id.clone(),
256                error: error.to_string(),
257            },
258        );
259    }
260}
261
262fn next_compaction_request_id() -> CompactionRequestId {
263    let ordinal = NEXT_COMPACTION_REQUEST.fetch_add(1, Ordering::Relaxed);
264    format!("compaction_request:{ordinal}")
265}
266
267fn write_json_payload_best_effort(
268    writer: &TraceWriter,
269    kind: RawPayloadKind,
270    payload: &impl Serialize,
271) -> Option<crate::RawPayloadRef> {
272    writer.write_json_payload(kind, payload).ok()
273}
274
275fn append_with_context_best_effort(
276    context: &EnabledCompactionTraceContext,
277    payload: RawTraceEventPayload,
278) {
279    let event_context = RawTraceEventContext {
280        thread_id: Some(context.thread_id.clone()),
281        codex_turn_id: Some(context.codex_turn_id.clone()),
282    };
283    let _ = context.writer.append_with_context(event_context, payload);
284}
codex_rollout_trace/compaction.rs

codex_rollout_trace/
compaction.rs
