From a4a812d27e3eb651720b43cdcde07ec7331f0a0d Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Mon, 13 Apr 2026 22:24:16 -0700 Subject: [PATCH] codex: remove realtime transcript correlation state --- .../endpoint/realtime_websocket/methods.rs | 329 +----------------- .../realtime_websocket/protocol_common.rs | 28 +- .../realtime_websocket/protocol_v2.rs | 7 +- codex-rs/protocol/src/protocol.rs | 30 +- 4 files changed, 18 insertions(+), 376 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index 4f2ec9c54c..14fc775b6a 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -18,7 +18,6 @@ use crate::provider::Provider; use codex_client::backoff; use codex_client::maybe_build_rustls_client_config_with_custom_ca; use codex_protocol::protocol::RealtimeTranscriptDelta; -use codex_protocol::protocol::RealtimeTranscriptDone; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use futures::SinkExt; use futures::StreamExt; @@ -213,26 +212,9 @@ pub struct RealtimeWebsocketEvents { is_closed: Arc, } -#[derive(Debug, Default, PartialEq, Eq)] +#[derive(Default)] struct ActiveTranscriptState { entries: Vec, - in_progress_parts: Vec, -} - -#[derive(Debug, PartialEq, Eq)] -struct ActiveTranscriptPart { - key: TranscriptPartKey, - entry_index: usize, - start: usize, - end: usize, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct TranscriptPartKey { - role: String, - item_id: Option, - output_index: Option, - content_index: Option, } impl RealtimeWebsocketConnection { @@ -432,25 +414,20 @@ impl RealtimeWebsocketEvents { let mut active_transcript = self.active_transcript.lock().await; match event { RealtimeEvent::InputAudioSpeechStarted(_) => {} - RealtimeEvent::InputTranscriptDelta(update) => { - append_transcript_delta(&mut active_transcript, "user", update); + RealtimeEvent::InputTranscriptDelta(RealtimeTranscriptDelta { delta, .. }) => { + append_transcript_delta(&mut active_transcript.entries, "user", delta); } - RealtimeEvent::InputTranscriptDone(update) => { - complete_transcript_entry(&mut active_transcript, "user", update); - } - RealtimeEvent::OutputTranscriptDelta(update) => { - append_transcript_delta(&mut active_transcript, "assistant", update); - } - RealtimeEvent::OutputTranscriptDone(update) => { - complete_transcript_entry(&mut active_transcript, "assistant", update); + RealtimeEvent::OutputTranscriptDelta(RealtimeTranscriptDelta { delta, .. }) => { + append_transcript_delta(&mut active_transcript.entries, "assistant", delta); } RealtimeEvent::HandoffRequested(handoff) => { if self.event_parser == RealtimeEventParser::V1 { handoff.active_transcript = std::mem::take(&mut active_transcript.entries); - active_transcript.in_progress_parts.clear(); } } RealtimeEvent::SessionUpdated { .. } + | RealtimeEvent::InputTranscriptDone(_) + | RealtimeEvent::OutputTranscriptDone(_) | RealtimeEvent::AudioOut(_) | RealtimeEvent::ResponseCreated(_) | RealtimeEvent::ResponseCancelled(_) @@ -462,182 +439,24 @@ impl RealtimeWebsocketEvents { } } -fn append_transcript_delta( - state: &mut ActiveTranscriptState, - role: &str, - update: &RealtimeTranscriptDelta, -) { - if update.delta.is_empty() { +fn append_transcript_delta(entries: &mut Vec, role: &str, delta: &str) { + if delta.is_empty() { return; } - let key = transcript_part_key( - role, - update.item_id.as_deref(), - update.output_index, - update.content_index, - ); - let key_has_metadata = transcript_part_key_has_metadata(&key); - - if let Some(part_index) = state - .in_progress_parts - .iter() - .position(|part| part.key == key) - { - let entry_index = state.in_progress_parts[part_index].entry_index; - if state - .entries - .get(entry_index) - .is_some_and(|entry| entry.role == role) - && (key_has_metadata || entry_index == state.entries.len().saturating_sub(1)) - { - let insert_at = state.in_progress_parts[part_index].end; - state.entries[entry_index] - .text - .insert_str(insert_at, &update.delta); - adjust_transcript_part_offsets( - state, - entry_index, - insert_at, - isize::try_from(update.delta.len()).unwrap_or(isize::MAX), - ); - state.in_progress_parts[part_index].end += update.delta.len(); - return; - } - - state.in_progress_parts.swap_remove(part_index); - } - - let entry_index = transcript_entry_index(&mut state.entries, role); - let start = state.entries[entry_index].text.len(); - state.entries[entry_index].text.push_str(&update.delta); - - state.in_progress_parts.push(ActiveTranscriptPart { - key, - entry_index, - start, - end: start + update.delta.len(), - }); -} - -fn complete_transcript_entry( - state: &mut ActiveTranscriptState, - role: &str, - update: &RealtimeTranscriptDone, -) { - let key = transcript_part_key( - role, - update.item_id.as_deref(), - update.output_index, - update.content_index, - ); - let key_has_metadata = transcript_part_key_has_metadata(&key); - let part_index = state - .in_progress_parts - .iter() - .position(|part| part.key == key); - - if let Some(part_index) = part_index { - let part = state.in_progress_parts.swap_remove(part_index); - let old_len = part.end.saturating_sub(part.start); - let entry_is_current = - key_has_metadata || part.entry_index == state.entries.len().saturating_sub(1); - let replaced = if entry_is_current - && let Some(entry) = state.entries.get_mut(part.entry_index) - && entry.role == role - && part.end <= entry.text.len() - { - // Done events carry the complete current part, not the complete - // message, so only replace the range accumulated for that part. - entry.text.replace_range(part.start..part.end, &update.text); - true - } else { - false - }; - - if replaced { - let length_delta = isize::try_from(update.text.len()).unwrap_or(isize::MAX) - - isize::try_from(old_len).unwrap_or(isize::MAX); - adjust_transcript_part_offsets(state, part.entry_index, part.end, length_delta); - return; - } - } - - let text = &update.text; - if text.is_empty() { - return; - } - - if let Some(last_entry) = state.entries.last_mut() + if let Some(last_entry) = entries.last_mut() && last_entry.role == role { - last_entry.text.push_str(text); + last_entry.text.push_str(delta); return; } - state.entries.push(RealtimeTranscriptEntry { + entries.push(RealtimeTranscriptEntry { role: role.to_string(), - text: text.to_string(), + text: delta.to_string(), }); } -fn transcript_part_key( - role: &str, - item_id: Option<&str>, - output_index: Option, - content_index: Option, -) -> TranscriptPartKey { - TranscriptPartKey { - role: role.to_string(), - item_id: item_id.map(str::to_string), - output_index, - content_index, - } -} - -fn transcript_part_key_has_metadata(key: &TranscriptPartKey) -> bool { - key.item_id.is_some() || key.output_index.is_some() || key.content_index.is_some() -} - -fn transcript_entry_index(entries: &mut Vec, role: &str) -> usize { - match entries.last() { - Some(last_entry) if last_entry.role == role => entries.len() - 1, - _ => { - entries.push(RealtimeTranscriptEntry { - role: role.to_string(), - text: String::new(), - }); - entries.len() - 1 - } - } -} - -fn adjust_transcript_part_offsets( - state: &mut ActiveTranscriptState, - entry_index: usize, - after: usize, - amount: isize, -) { - if amount == 0 { - return; - } - - for part in &mut state.in_progress_parts { - if part.entry_index == entry_index && part.start >= after { - adjust_offset(&mut part.start, amount); - adjust_offset(&mut part.end, amount); - } - } -} - -fn adjust_offset(offset: &mut usize, amount: isize) { - if amount.is_positive() { - *offset += amount.unsigned_abs(); - } else { - *offset -= amount.unsigned_abs(); - } -} - pub struct RealtimeWebsocketClient { provider: Provider, } @@ -1039,7 +858,6 @@ mod tests { Some(RealtimeEvent::InputTranscriptDelta( RealtimeTranscriptDelta { delta: "hello ".to_string(), - ..Default::default() } )) ); @@ -1058,7 +876,6 @@ mod tests { Some(RealtimeEvent::OutputTranscriptDelta( RealtimeTranscriptDelta { delta: "hi".to_string(), - ..Default::default() } )) ); @@ -1104,9 +921,6 @@ mod tests { Some(RealtimeEvent::InputTranscriptDelta( RealtimeTranscriptDelta { delta: "hello".to_string(), - item_id: Some("item_input_1".to_string()), - output_index: None, - content_index: Some(0), } )) ); @@ -1128,9 +942,6 @@ mod tests { Some(RealtimeEvent::OutputTranscriptDone( RealtimeTranscriptDone { text: "all done".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: Some(0), - content_index: Some(1), } )) ); @@ -1157,122 +968,11 @@ mod tests { Some(RealtimeEvent::OutputTranscriptDone( RealtimeTranscriptDone { text: "hello world".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: None, - content_index: None, } )) ); } - #[test] - fn complete_transcript_entry_replaces_current_part_only() { - let mut state = ActiveTranscriptState::default(); - let first_part_delta = RealtimeTranscriptDelta { - delta: "hello".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: Some(0), - content_index: Some(0), - }; - let second_part_delta = RealtimeTranscriptDelta { - delta: "beta".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: Some(0), - content_index: Some(1), - }; - let first_part_done = RealtimeTranscriptDone { - text: "hello!".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: Some(0), - content_index: Some(0), - }; - let second_part_done = RealtimeTranscriptDone { - text: " beta".to_string(), - item_id: Some("item_output_1".to_string()), - output_index: Some(0), - content_index: Some(1), - }; - - append_transcript_delta(&mut state, "assistant", &first_part_delta); - append_transcript_delta(&mut state, "assistant", &second_part_delta); - complete_transcript_entry(&mut state, "assistant", &first_part_done); - complete_transcript_entry(&mut state, "assistant", &second_part_done); - - assert_eq!( - state, - ActiveTranscriptState { - entries: vec![RealtimeTranscriptEntry { - role: "assistant".to_string(), - text: "hello! beta".to_string(), - }], - in_progress_parts: Vec::new(), - } - ); - } - - #[test] - fn unkeyed_transcript_deltas_start_new_entry_after_role_change() { - let mut state = ActiveTranscriptState::default(); - let assistant_context = RealtimeTranscriptDelta { - delta: "assistant context".to_string(), - ..Default::default() - }; - let delegated_query = RealtimeTranscriptDelta { - delta: "delegated query".to_string(), - ..Default::default() - }; - let assist_confirm = RealtimeTranscriptDelta { - delta: "assist confirm".to_string(), - ..Default::default() - }; - - append_transcript_delta(&mut state, "assistant", &assistant_context); - append_transcript_delta(&mut state, "user", &delegated_query); - append_transcript_delta(&mut state, "assistant", &assist_confirm); - - assert_eq!( - state, - ActiveTranscriptState { - entries: vec![ - RealtimeTranscriptEntry { - role: "assistant".to_string(), - text: "assistant context".to_string(), - }, - RealtimeTranscriptEntry { - role: "user".to_string(), - text: "delegated query".to_string(), - }, - RealtimeTranscriptEntry { - role: "assistant".to_string(), - text: "assist confirm".to_string(), - }, - ], - in_progress_parts: vec![ - ActiveTranscriptPart { - key: transcript_part_key( - "user", /*item_id*/ None, /*output_index*/ None, - /*content_index*/ None, - ), - entry_index: 1, - start: 0, - end: "delegated query".len(), - }, - ActiveTranscriptPart { - key: transcript_part_key( - "assistant", - /*item_id*/ None, - /*output_index*/ None, - /*content_index*/ None, - ), - entry_index: 2, - start: 0, - end: "assist confirm".len(), - }, - ], - } - ); - } - #[test] fn parse_realtime_v2_output_audio_delta_defaults_audio_shape() { let payload = json!({ @@ -1808,7 +1508,6 @@ mod tests { input_delta_event, RealtimeEvent::InputTranscriptDelta(RealtimeTranscriptDelta { delta: "delegate ".to_string(), - ..Default::default() }) ); @@ -1821,7 +1520,6 @@ mod tests { input_delta_event, RealtimeEvent::InputTranscriptDelta(RealtimeTranscriptDelta { delta: "now".to_string(), - ..Default::default() }) ); @@ -1834,7 +1532,6 @@ mod tests { output_delta_event, RealtimeEvent::OutputTranscriptDelta(RealtimeTranscriptDelta { delta: "working".to_string(), - ..Default::default() }) ); diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_common.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_common.rs index 00513dc84e..c89c5ea4d0 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_common.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_common.rs @@ -51,12 +51,7 @@ pub(super) fn parse_transcript_delta_event( .get(field) .and_then(Value::as_str) .map(str::to_string) - .map(|delta| RealtimeTranscriptDelta { - delta, - item_id: parse_string_field(parsed, "item_id"), - output_index: parse_u32_field(parsed, "output_index"), - content_index: parse_u32_field(parsed, "content_index"), - }) + .map(|delta| RealtimeTranscriptDelta { delta }) } pub(super) fn parse_transcript_done_event( @@ -67,26 +62,7 @@ pub(super) fn parse_transcript_done_event( .get(field) .and_then(Value::as_str) .map(str::to_string) - .map(|text| RealtimeTranscriptDone { - text, - item_id: parse_string_field(parsed, "item_id"), - output_index: parse_u32_field(parsed, "output_index"), - content_index: parse_u32_field(parsed, "content_index"), - }) -} - -fn parse_string_field(parsed: &Value, field: &str) -> Option { - parsed - .get(field) - .and_then(Value::as_str) - .map(str::to_string) -} - -fn parse_u32_field(parsed: &Value, field: &str) -> Option { - parsed - .get(field) - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) + .map(|text| RealtimeTranscriptDone { text }) } pub(super) fn parse_error_event(parsed: &Value) -> Option { diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_v2.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_v2.rs index 4cea1535d6..7c1c8d8b87 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_v2.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol_v2.rs @@ -151,12 +151,7 @@ fn parse_item_done_transcript(item: &JsonMap) -> Option Some(RealtimeEvent::InputTranscriptDone(done)), "assistant" => Some(RealtimeEvent::OutputTranscriptDone(done)), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 57e17a287b..7431e5f123 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -294,40 +294,14 @@ pub struct RealtimeAudioFrame { pub item_id: Option, } -#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] pub struct RealtimeTranscriptDelta { pub delta: String, - /// Internal key material used to reconcile streamed transcript parts. - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub item_id: Option, - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub output_index: Option, - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub content_index: Option, } -#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] pub struct RealtimeTranscriptDone { pub text: String, - /// Internal key material used to reconcile streamed transcript parts. - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub item_id: Option, - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub output_index: Option, - #[serde(skip)] - #[schemars(skip)] - #[ts(skip)] - pub content_index: Option, } #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]