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 dcd330322e..9b7956b346 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -211,9 +211,17 @@ pub struct RealtimeWebsocketEvents { is_closed: Arc, } -#[derive(Default)] +#[derive(Debug, Default, PartialEq, Eq)] struct ActiveTranscriptState { entries: Vec, + in_progress_parts: Vec, +} + +#[derive(Debug, PartialEq, Eq)] +struct ActiveTranscriptPart { + role: String, + entry_index: usize, + start: usize, } impl RealtimeWebsocketConnection { @@ -414,24 +422,21 @@ impl RealtimeWebsocketEvents { match event { RealtimeEvent::InputAudioSpeechStarted(_) => {} RealtimeEvent::InputTranscriptDelta(update) => { - append_transcript_delta(&mut active_transcript.entries, "user", &update.delta); + append_transcript_delta(&mut active_transcript, "user", &update.delta); } RealtimeEvent::InputTranscriptDone(update) => { - complete_transcript_entry(&mut active_transcript.entries, "user", &update.text); + complete_transcript_entry(&mut active_transcript, "user", &update.text); } RealtimeEvent::OutputTranscriptDelta(update) => { - append_transcript_delta(&mut active_transcript.entries, "assistant", &update.delta); + append_transcript_delta(&mut active_transcript, "assistant", &update.delta); } RealtimeEvent::OutputTranscriptDone(update) => { - complete_transcript_entry( - &mut active_transcript.entries, - "assistant", - &update.text, - ); + complete_transcript_entry(&mut active_transcript, "assistant", &update.text); } 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 { .. } @@ -446,39 +451,80 @@ impl RealtimeWebsocketEvents { } } -fn append_transcript_delta(entries: &mut Vec, role: &str, delta: &str) { +fn append_transcript_delta(state: &mut ActiveTranscriptState, role: &str, delta: &str) { if delta.is_empty() { return; } - if let Some(last_entry) = entries.last_mut() - && last_entry.role == role - { - last_entry.text.push_str(delta); - return; + let entry_index = match state.entries.last() { + Some(last_entry) if last_entry.role == role => state.entries.len() - 1, + _ => { + state.entries.push(RealtimeTranscriptEntry { + role: role.to_string(), + text: String::new(), + }); + state.entries.len() - 1 + } + }; + + let part_index = state + .in_progress_parts + .iter() + .position(|part| part.role == role); + let start = part_index + .and_then(|index| state.in_progress_parts.get(index)) + .filter(|part| part.entry_index == entry_index) + .map(|part| part.start) + .unwrap_or_else(|| state.entries[entry_index].text.len()); + + let active_part = ActiveTranscriptPart { + role: role.to_string(), + entry_index, + start, + }; + if let Some(index) = part_index { + state.in_progress_parts[index] = active_part; + } else { + state.in_progress_parts.push(active_part); } - entries.push(RealtimeTranscriptEntry { - role: role.to_string(), - text: delta.to_string(), - }); + state.entries[entry_index].text.push_str(delta); } -fn complete_transcript_entry(entries: &mut Vec, role: &str, text: &str) { +fn complete_transcript_entry(state: &mut ActiveTranscriptState, role: &str, text: &str) { + let part_index = state + .in_progress_parts + .iter() + .position(|part| part.role == role); + if text.is_empty() { + if let Some(part_index) = part_index { + state.in_progress_parts.swap_remove(part_index); + } return; } - // Final transcript events carry the complete part. Replace the in-progress entry - // instead of appending, so active transcript state stays usable after deltas. - if let Some(last_entry) = entries.last_mut() + if let Some(part_index) = part_index { + let part = state.in_progress_parts.swap_remove(part_index); + if let Some(entry) = state.entries.get_mut(part.entry_index) + && entry.role == role + && part.start <= entry.text.len() + { + // Done events carry the complete current part, not the complete + // message, so only replace the suffix accumulated for that part. + entry.text.replace_range(part.start.., text); + return; + } + } + + if let Some(last_entry) = state.entries.last_mut() && last_entry.role == role { - last_entry.text = text.to_string(); + last_entry.text.push_str(text); return; } - entries.push(RealtimeTranscriptEntry { + state.entries.push(RealtimeTranscriptEntry { role: role.to_string(), text: text.to_string(), }); @@ -969,6 +1015,27 @@ mod tests { ); } + #[test] + fn complete_transcript_entry_replaces_current_part_only() { + let mut state = ActiveTranscriptState::default(); + + append_transcript_delta(&mut state, "assistant", "hello"); + complete_transcript_entry(&mut state, "assistant", "hello"); + append_transcript_delta(&mut state, "assistant", " wor"); + complete_transcript_entry(&mut state, "assistant", " world"); + + assert_eq!( + state, + ActiveTranscriptState { + entries: vec![RealtimeTranscriptEntry { + role: "assistant".to_string(), + text: "hello world".to_string(), + }], + in_progress_parts: Vec::new(), + } + ); + } + #[test] fn parse_realtime_v2_output_audio_delta_defaults_audio_shape() { let payload = json!({