mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Wire realtime api to core (#12268)
- Introduce `RealtimeConversationManager` for realtime API management - Add `op::conversation` to start conversation, insert audio, insert text, and close conversation. - emit conversation lifecycle and realtime events. - Move shared realtime payload types into codex-protocol and add core e2e websocket tests for start/replace/transport-close paths. Things to consider: - Should we use the same `op::` and `Events` channel to carry audio? I think we should try this simple approach and later we can create separate one if the channels got congested. - Sending text updates to the client: we can start simple and later restrict that. - Provider auth isn't wired for now intentionally
This commit is contained in:
Ahmed Ibrahim
committed by
GitHub
Unverified
GitHub
parent
936e744c93
commit
6817f0be8a
@@ -30,6 +30,11 @@ use crate::features::maybe_push_unstable_features_warning;
|
||||
use crate::models_manager::manager::ModelsManager;
|
||||
use crate::parse_command::parse_command;
|
||||
use crate::parse_turn_item;
|
||||
use crate::realtime_conversation::RealtimeConversationManager;
|
||||
use crate::realtime_conversation::handle_audio as handle_realtime_conversation_audio;
|
||||
use crate::realtime_conversation::handle_close as handle_realtime_conversation_close;
|
||||
use crate::realtime_conversation::handle_start as handle_realtime_conversation_start;
|
||||
use crate::realtime_conversation::handle_text as handle_realtime_conversation_text;
|
||||
use crate::rollout::session_index;
|
||||
use crate::stream_events_utils::HandleOutputCtx;
|
||||
use crate::stream_events_utils::handle_non_tool_response_item;
|
||||
@@ -282,7 +287,7 @@ pub struct CodexSpawnOk {
|
||||
}
|
||||
|
||||
pub(crate) const INITIAL_SUBMIT_ID: &str = "";
|
||||
pub(crate) const SUBMISSION_CHANNEL_CAPACITY: usize = 64;
|
||||
pub(crate) const SUBMISSION_CHANNEL_CAPACITY: usize = 512;
|
||||
const CYBER_VERIFY_URL: &str = "https://chatgpt.com/cyber";
|
||||
const CYBER_SAFETY_URL: &str = "https://developers.openai.com/codex/concepts/cyber-safety";
|
||||
|
||||
@@ -523,11 +528,13 @@ pub(crate) struct Session {
|
||||
/// session.
|
||||
features: Features,
|
||||
pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
|
||||
pub(crate) conversation: Arc<RealtimeConversationManager>,
|
||||
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
|
||||
pub(crate) services: SessionServices,
|
||||
js_repl: Arc<JsReplHandle>,
|
||||
next_internal_sub_id: AtomicU64,
|
||||
}
|
||||
|
||||
/// The context needed for a single turn of the thread.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct TurnContext {
|
||||
@@ -1357,6 +1364,7 @@ impl Session {
|
||||
state: Mutex::new(state),
|
||||
features: config.features.clone(),
|
||||
pending_mcp_server_refresh_config: Mutex::new(None),
|
||||
conversation: Arc::new(RealtimeConversationManager::new()),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
js_repl,
|
||||
@@ -1961,7 +1969,7 @@ impl Session {
|
||||
state.take_startup_regular_task()
|
||||
}
|
||||
|
||||
async fn get_config(&self) -> std::sync::Arc<Config> {
|
||||
pub(crate) async fn get_config(&self) -> std::sync::Arc<Config> {
|
||||
let state = self.state.lock().await;
|
||||
state
|
||||
.session_configuration
|
||||
@@ -1969,6 +1977,11 @@ impl Session {
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn provider(&self) -> ModelProviderInfo {
|
||||
let state = self.state.lock().await;
|
||||
state.session_configuration.provider.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn reload_user_config_layer(&self) {
|
||||
let config_toml_path = {
|
||||
let state = self.state.lock().await;
|
||||
@@ -3170,6 +3183,29 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
|
||||
Op::CleanBackgroundTerminals => {
|
||||
handlers::clean_background_terminals(&sess).await;
|
||||
}
|
||||
Op::RealtimeConversationStart(params) => {
|
||||
if let Err(err) =
|
||||
handle_realtime_conversation_start(&sess, sub.id.clone(), params).await
|
||||
{
|
||||
sess.send_event_raw(Event {
|
||||
id: sub.id.clone(),
|
||||
msg: EventMsg::Error(ErrorEvent {
|
||||
message: err.to_string(),
|
||||
codex_error_info: Some(CodexErrorInfo::Other),
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Op::RealtimeConversationAudio(params) => {
|
||||
handle_realtime_conversation_audio(&sess, sub.id.clone(), params).await;
|
||||
}
|
||||
Op::RealtimeConversationText(params) => {
|
||||
handle_realtime_conversation_text(&sess, sub.id.clone(), params).await;
|
||||
}
|
||||
Op::RealtimeConversationClose => {
|
||||
handle_realtime_conversation_close(&sess, sub.id.clone()).await;
|
||||
}
|
||||
Op::OverrideTurnContext {
|
||||
cwd,
|
||||
approval_policy,
|
||||
@@ -4015,6 +4051,7 @@ mod handlers {
|
||||
|
||||
pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
|
||||
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
|
||||
let _ = sess.conversation.shutdown().await;
|
||||
sess.services
|
||||
.unified_exec_manager
|
||||
.terminate_all_processes()
|
||||
@@ -7467,6 +7504,7 @@ mod tests {
|
||||
state: Mutex::new(state),
|
||||
features: config.features.clone(),
|
||||
pending_mcp_server_refresh_config: Mutex::new(None),
|
||||
conversation: Arc::new(RealtimeConversationManager::new()),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
js_repl,
|
||||
@@ -7623,6 +7661,7 @@ mod tests {
|
||||
state: Mutex::new(state),
|
||||
features: config.features.clone(),
|
||||
pending_mcp_server_refresh_config: Mutex::new(None),
|
||||
conversation: Arc::new(RealtimeConversationManager::new()),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
js_repl,
|
||||
|
||||
@@ -13,6 +13,7 @@ pub mod auth;
|
||||
mod client;
|
||||
mod client_common;
|
||||
pub mod codex;
|
||||
mod realtime_conversation;
|
||||
pub use codex::SteerInputError;
|
||||
mod codex_thread;
|
||||
mod compact_remote;
|
||||
|
||||
@@ -0,0 +1,354 @@
|
||||
use crate::CodexAuth;
|
||||
use crate::api_bridge::map_api_error;
|
||||
use crate::codex::Session;
|
||||
use crate::default_client::default_headers;
|
||||
use crate::error::CodexErr;
|
||||
use crate::error::Result as CodexResult;
|
||||
use async_channel::Receiver;
|
||||
use async_channel::Sender;
|
||||
use async_channel::TrySendError;
|
||||
use codex_api::Provider as ApiProvider;
|
||||
use codex_api::RealtimeAudioFrame;
|
||||
use codex_api::RealtimeEvent;
|
||||
use codex_api::RealtimeSessionConfig;
|
||||
use codex_api::RealtimeWebsocketClient;
|
||||
use codex_api::endpoint::realtime_websocket::RealtimeWebsocketEvents;
|
||||
use codex_api::endpoint::realtime_websocket::RealtimeWebsocketWriter;
|
||||
use codex_protocol::protocol::CodexErrorInfo;
|
||||
use codex_protocol::protocol::ConversationAudioParams;
|
||||
use codex_protocol::protocol::ConversationStartParams;
|
||||
use codex_protocol::protocol::ConversationTextParams;
|
||||
use codex_protocol::protocol::ErrorEvent;
|
||||
use codex_protocol::protocol::Event;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::RealtimeConversationClosedEvent;
|
||||
use codex_protocol::protocol::RealtimeConversationRealtimeEvent;
|
||||
use codex_protocol::protocol::RealtimeConversationStartedEvent;
|
||||
use http::HeaderMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::error;
|
||||
use tracing::warn;
|
||||
|
||||
const AUDIO_IN_QUEUE_CAPACITY: usize = 256;
|
||||
const TEXT_IN_QUEUE_CAPACITY: usize = 64;
|
||||
const OUTPUT_EVENTS_QUEUE_CAPACITY: usize = 256;
|
||||
|
||||
pub(crate) struct RealtimeConversationManager {
|
||||
state: Mutex<Option<ConversationState>>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
struct ConversationState {
|
||||
audio_tx: Sender<RealtimeAudioFrame>,
|
||||
text_tx: Sender<String>,
|
||||
task: JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl RealtimeConversationManager {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
state: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn running_state(&self) -> Option<()> {
|
||||
let state = self.state.lock().await;
|
||||
state.as_ref().map(|_| ())
|
||||
}
|
||||
|
||||
pub(crate) async fn start(
|
||||
&self,
|
||||
api_provider: ApiProvider,
|
||||
extra_headers: Option<HeaderMap>,
|
||||
prompt: String,
|
||||
session_id: Option<String>,
|
||||
) -> CodexResult<Receiver<RealtimeEvent>> {
|
||||
let previous_state = {
|
||||
let mut guard = self.state.lock().await;
|
||||
guard.take()
|
||||
};
|
||||
if let Some(state) = previous_state {
|
||||
state.task.abort();
|
||||
let _ = state.task.await;
|
||||
}
|
||||
|
||||
let session_config = RealtimeSessionConfig { prompt, session_id };
|
||||
let client = RealtimeWebsocketClient::new(api_provider);
|
||||
let connection = client
|
||||
.connect(
|
||||
session_config,
|
||||
extra_headers.unwrap_or_default(),
|
||||
default_headers(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_api_error)?;
|
||||
|
||||
let writer = connection.writer();
|
||||
let events = connection.events();
|
||||
let (audio_tx, audio_rx) =
|
||||
async_channel::bounded::<RealtimeAudioFrame>(AUDIO_IN_QUEUE_CAPACITY);
|
||||
let (text_tx, text_rx) = async_channel::bounded::<String>(TEXT_IN_QUEUE_CAPACITY);
|
||||
let (events_tx, events_rx) =
|
||||
async_channel::bounded::<RealtimeEvent>(OUTPUT_EVENTS_QUEUE_CAPACITY);
|
||||
|
||||
let task = spawn_realtime_input_task(writer, events, text_rx, audio_rx, events_tx);
|
||||
|
||||
let mut guard = self.state.lock().await;
|
||||
*guard = Some(ConversationState {
|
||||
audio_tx,
|
||||
text_tx,
|
||||
task,
|
||||
});
|
||||
Ok(events_rx)
|
||||
}
|
||||
|
||||
pub(crate) async fn audio_in(&self, frame: RealtimeAudioFrame) -> CodexResult<()> {
|
||||
let sender = {
|
||||
let guard = self.state.lock().await;
|
||||
guard.as_ref().map(|state| state.audio_tx.clone())
|
||||
};
|
||||
|
||||
let Some(sender) = sender else {
|
||||
return Err(CodexErr::InvalidRequest(
|
||||
"conversation is not running".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
match sender.try_send(frame) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(TrySendError::Full(_)) => {
|
||||
warn!("dropping input audio frame due to full queue");
|
||||
Ok(())
|
||||
}
|
||||
Err(TrySendError::Closed(_)) => Err(CodexErr::InvalidRequest(
|
||||
"conversation is not running".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn text_in(&self, text: String) -> CodexResult<()> {
|
||||
let sender = {
|
||||
let guard = self.state.lock().await;
|
||||
guard.as_ref().map(|state| state.text_tx.clone())
|
||||
};
|
||||
|
||||
let Some(sender) = sender else {
|
||||
return Err(CodexErr::InvalidRequest(
|
||||
"conversation is not running".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
sender
|
||||
.send(text)
|
||||
.await
|
||||
.map_err(|_| CodexErr::InvalidRequest("conversation is not running".to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&self) -> CodexResult<()> {
|
||||
let state = {
|
||||
let mut guard = self.state.lock().await;
|
||||
guard.take()
|
||||
};
|
||||
|
||||
if let Some(state) = state {
|
||||
state.task.abort();
|
||||
let _ = state.task.await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_start(
|
||||
sess: &Arc<Session>,
|
||||
sub_id: String,
|
||||
params: ConversationStartParams,
|
||||
) -> CodexResult<()> {
|
||||
let provider = sess.provider().await;
|
||||
let auth = sess.services.auth_manager.auth().await;
|
||||
let api_provider = provider.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
|
||||
|
||||
let requested_session_id = params
|
||||
.session_id
|
||||
.or_else(|| Some(sess.conversation_id.to_string()));
|
||||
let events_rx = match sess
|
||||
.conversation
|
||||
.start(
|
||||
api_provider,
|
||||
None,
|
||||
params.prompt,
|
||||
requested_session_id.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(events_rx) => events_rx,
|
||||
Err(err) => {
|
||||
send_conversation_error(sess, sub_id, err.to_string(), CodexErrorInfo::Other).await;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
sess.send_event_raw(Event {
|
||||
id: sub_id.clone(),
|
||||
msg: EventMsg::RealtimeConversationStarted(RealtimeConversationStartedEvent {
|
||||
session_id: requested_session_id,
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
|
||||
let sess_clone = Arc::clone(sess);
|
||||
tokio::spawn(async move {
|
||||
let ev = |msg| Event {
|
||||
id: sub_id.clone(),
|
||||
msg,
|
||||
};
|
||||
while let Ok(event) = events_rx.recv().await {
|
||||
sess_clone
|
||||
.send_event_raw(ev(EventMsg::RealtimeConversationRealtime(
|
||||
RealtimeConversationRealtimeEvent { payload: event },
|
||||
)))
|
||||
.await;
|
||||
}
|
||||
if let Some(()) = sess_clone.conversation.running_state().await {
|
||||
sess_clone
|
||||
.send_event_raw(ev(EventMsg::RealtimeConversationClosed(
|
||||
RealtimeConversationClosedEvent {
|
||||
reason: Some("transport_closed".to_string()),
|
||||
},
|
||||
)))
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_audio(
|
||||
sess: &Arc<Session>,
|
||||
sub_id: String,
|
||||
params: ConversationAudioParams,
|
||||
) {
|
||||
if let Err(err) = sess.conversation.audio_in(params.frame).await {
|
||||
send_conversation_error(sess, sub_id, err.to_string(), CodexErrorInfo::BadRequest).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_text(
|
||||
sess: &Arc<Session>,
|
||||
sub_id: String,
|
||||
params: ConversationTextParams,
|
||||
) {
|
||||
if let Err(err) = sess.conversation.text_in(params.text).await {
|
||||
send_conversation_error(sess, sub_id, err.to_string(), CodexErrorInfo::BadRequest).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_close(sess: &Arc<Session>, sub_id: String) {
|
||||
match sess.conversation.shutdown().await {
|
||||
Ok(()) => {
|
||||
sess.send_event_raw(Event {
|
||||
id: sub_id,
|
||||
msg: EventMsg::RealtimeConversationClosed(RealtimeConversationClosedEvent {
|
||||
reason: Some("requested".to_string()),
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(err) => {
|
||||
send_conversation_error(sess, sub_id, err.to_string(), CodexErrorInfo::Other).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_realtime_input_task(
|
||||
writer: RealtimeWebsocketWriter,
|
||||
events: RealtimeWebsocketEvents,
|
||||
text_rx: Receiver<String>,
|
||||
audio_rx: Receiver<RealtimeAudioFrame>,
|
||||
events_tx: Sender<RealtimeEvent>,
|
||||
) -> JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
text = text_rx.recv() => {
|
||||
match text {
|
||||
Ok(text) => {
|
||||
if let Err(err) = writer.send_conversation_item_create(text).await {
|
||||
let mapped_error = map_api_error(err);
|
||||
warn!("failed to send input text: {mapped_error}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
event = events.next_event() => {
|
||||
match event {
|
||||
Ok(Some(event)) => {
|
||||
let should_stop = matches!(&event, RealtimeEvent::Error(_));
|
||||
if events_tx.send(event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if should_stop {
|
||||
error!("realtime stream error event received");
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
let _ = events_tx
|
||||
.send(RealtimeEvent::Error(
|
||||
"realtime websocket connection is closed".to_string(),
|
||||
))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Err(err) => {
|
||||
let mapped_error = map_api_error(err);
|
||||
if events_tx
|
||||
.send(RealtimeEvent::Error(mapped_error.to_string()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
error!("realtime stream closed: {mapped_error}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
frame = audio_rx.recv() => {
|
||||
match frame {
|
||||
Ok(frame) => {
|
||||
if let Err(err) = writer.send_audio_frame(frame).await {
|
||||
let mapped_error = map_api_error(err);
|
||||
error!("failed to send input audio: {mapped_error}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn send_conversation_error(
|
||||
sess: &Arc<Session>,
|
||||
sub_id: String,
|
||||
message: String,
|
||||
codex_error_info: CodexErrorInfo,
|
||||
) {
|
||||
sess.send_event_raw(Event {
|
||||
id: sub_id,
|
||||
msg: EventMsg::Error(ErrorEvent {
|
||||
message,
|
||||
codex_error_info: Some(codex_error_info),
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -122,6 +122,9 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
|
||||
| EventMsg::CollabCloseEnd(_)
|
||||
| EventMsg::CollabResumeEnd(_) => Some(EventPersistenceMode::Extended),
|
||||
EventMsg::Warning(_)
|
||||
| EventMsg::RealtimeConversationStarted(_)
|
||||
| EventMsg::RealtimeConversationRealtime(_)
|
||||
| EventMsg::RealtimeConversationClosed(_)
|
||||
| EventMsg::ModelReroute(_)
|
||||
| EventMsg::AgentMessageDelta(_)
|
||||
| EventMsg::AgentReasoningDelta(_)
|
||||
|
||||
Reference in New Issue
Block a user