valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use ConversationId instead of raw Uuids (#3282)

Browse Source 
We're trying to migrate from `session_id: Uuid` to `conversation_id:
ConversationId`. Not only does this give us more type safety but it
unifies our terminology across Codex and with the implementation of
session resuming, a conversation (which can span multiple sessions) is
more appropriate.

I started this impl on https://github.com/openai/codex/pull/3219 as part
of getting resume working in the extension but it's big enough that it
should be broken out.
This commit is contained in:
Gabriel Peal
2025-09-07 20:22:25 -07:00
committed by GitHub
parent 58d77ca4e7
commit c8fab51372
23 changed files with 213 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
codex-rs/mcp-server/src/codex_message_processor.rs
View File

@@ -99,7 +99,7 @@ pub(crate) struct CodexMessageProcessor {
conversation_listeners: HashMap<Uuid, oneshot::Sender<()>>,
active_login: Arc<Mutex<Option<ActiveLogin>>>,
// Queue of pending interrupt requests per conversation. We reply when TurnAborted arrives.
pending_interrupts: Arc<Mutex<HashMap<Uuid, Vec<RequestId>>>>,
pending_interrupts: Arc<Mutex<HashMap<ConversationId, Vec<RequestId>>>>,
}
impl CodexMessageProcessor {
@@ -511,7 +511,7 @@ impl CodexMessageProcessor {
..
} = conversation_id;
let response = NewConversationResponse {
conversation_id: ConversationId(conversation_id),
conversation_id,
model: session_configured.model,
};
self.outgoing.send_response(request_id, response).await;
@@ -632,7 +632,7 @@ impl CodexMessageProcessor {
// Reply with conversation id + model and initial messages (when present)
let response = codex_protocol::mcp_protocol::ResumeConversationResponse {
conversation_id: ConversationId(conversation_id),
conversation_id,
model: session_configured.model.clone(),
initial_messages: session_configured.initial_messages.clone(),
};
@@ -656,7 +656,7 @@ impl CodexMessageProcessor {
} = params;
let Ok(conversation) = self
.conversation_manager
.get_conversation(conversation_id.0)
.get_conversation(conversation_id)
.await
else {
let error = JSONRPCErrorError {
@@ -704,7 +704,7 @@ impl CodexMessageProcessor {
let Ok(conversation) = self
.conversation_manager
.get_conversation(conversation_id.0)
.get_conversation(conversation_id)
.await
else {
let error = JSONRPCErrorError {
@@ -750,7 +750,7 @@ impl CodexMessageProcessor {
let InterruptConversationParams { conversation_id } = params;
let Ok(conversation) = self
.conversation_manager
.get_conversation(conversation_id.0)
.get_conversation(conversation_id)
.await
else {
let error = JSONRPCErrorError {
@@ -765,7 +765,7 @@ impl CodexMessageProcessor {
// Record the pending interrupt so we can reply when TurnAborted arrives.
{
let mut map = self.pending_interrupts.lock().await;
map.entry(conversation_id.0).or_default().push(request_id);
map.entry(conversation_id).or_default().push(request_id);
}
// Submit the interrupt; we'll respond upon TurnAborted.
@@ -780,12 +780,12 @@ impl CodexMessageProcessor {
let AddConversationListenerParams { conversation_id } = params;
let Ok(conversation) = self
.conversation_manager
.get_conversation(conversation_id.0)
.get_conversation(conversation_id)
.await
else {
let error = JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("conversation not found: {}", conversation_id.0),
message: format!("conversation not found: {conversation_id}"),
data: None,
};
self.outgoing.send_error(request_id, error).await;
@@ -898,7 +898,7 @@ async fn apply_bespoke_event_handling(
conversation_id: ConversationId,
conversation: Arc<CodexConversation>,
outgoing: Arc<OutgoingMessageSender>,
pending_interrupts: Arc<Mutex<HashMap<Uuid, Vec<RequestId>>>>,
pending_interrupts: Arc<Mutex<HashMap<ConversationId, Vec<RequestId>>>>,
) {
let Event { id: event_id, msg } = event;
match msg {
@@ -951,7 +951,7 @@ async fn apply_bespoke_event_handling(
EventMsg::TurnAborted(turn_aborted_event) => {
let pending = {
let mut map = pending_interrupts.lock().await;
map.remove(&conversation_id.0).unwrap_or_default()
map.remove(&conversation_id).unwrap_or_default()
};
if !pending.is_empty() {
let response = InterruptConversationResponse {

19
codex-rs/mcp-server/src/codex_tool_config.rs
View File

@@ -181,8 +181,8 @@ impl CodexToolCallParam {
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct CodexToolCallReplyParam {
/// The *session id* for this conversation.
pub session_id: String,
/// The conversation id for this Codex session.
pub conversation_id: String,
/// The *next user prompt* to continue the Codex conversation.
pub prompt: String,
@@ -213,7 +213,8 @@ pub(crate) fn create_tool_for_codex_tool_call_reply_param() -> Tool {
input_schema: tool_input_schema,
output_schema: None,
description: Some(
"Continue a Codex session by providing the session id and prompt.".to_string(),
"Continue a Codex conversation by providing the conversation id and prompt."
.to_string(),
),
annotations: None,
}
@@ -308,21 +309,21 @@ mod tests {
let tool = create_tool_for_codex_tool_call_reply_param();
let tool_json = serde_json::to_value(&tool).expect("tool serializes");
let expected_tool_json = serde_json::json!({
"description": "Continue a Codex session by providing the session id and prompt.",
"description": "Continue a Codex conversation by providing the conversation id and prompt.",
"inputSchema": {
"properties": {
"conversationId": {
"description": "The conversation id for this Codex session.",
"type": "string"
},
"prompt": {
"description": "The *next user prompt* to continue the Codex conversation.",
"type": "string"
},
"sessionId": {
"description": "The *session id* for this conversation.",
"type": "string"
},
},
"required": [
"conversationId",
"prompt",
"sessionId",
],
"type": "object",
},

12
codex-rs/mcp-server/src/codex_tool_runner.rs
View File

@@ -18,13 +18,13 @@ use codex_core::protocol::InputItem;
use codex_core::protocol::Op;
use codex_core::protocol::Submission;
use codex_core::protocol::TaskCompleteEvent;
use codex_protocol::mcp_protocol::ConversationId;
use mcp_types::CallToolResult;
use mcp_types::ContentBlock;
use mcp_types::RequestId;
use mcp_types::TextContent;
use serde_json::json;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::exec_approval::handle_exec_approval_request;
use crate::outgoing_message::OutgoingMessageSender;
@@ -43,7 +43,7 @@ pub async fn run_codex_tool_session(
config: CodexConfig,
outgoing: Arc<OutgoingMessageSender>,
conversation_manager: Arc<ConversationManager>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
) {
let NewConversation {
conversation_id,
@@ -119,13 +119,13 @@ pub async fn run_codex_tool_session_reply(
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
prompt: String,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
session_id: Uuid,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
conversation_id: ConversationId,
) {
running_requests_id_to_codex_uuid
.lock()
.await
.insert(request_id.clone(), session_id);
.insert(request_id.clone(), conversation_id);
if let Err(e) = conversation
.submit(Op::UserInput {
items: vec![InputItem::Text { text: prompt }],
@@ -154,7 +154,7 @@ async fn run_codex_tool_session_inner(
codex: Arc<CodexConversation>,
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
) {
let request_id_str = match &request_id {
RequestId::String(s) => s.clone(),

48
codex-rs/mcp-server/src/message_processor.rs
View File

@@ -9,6 +9,7 @@ use crate::codex_tool_config::create_tool_for_codex_tool_call_reply_param;
use crate::error_code::INVALID_REQUEST_ERROR_CODE;
use crate::outgoing_message::OutgoingMessageSender;
use codex_protocol::mcp_protocol::ClientRequest;
use codex_protocol::mcp_protocol::ConversationId;
use codex_core::AuthManager;
use codex_core::ConversationManager;
@@ -41,7 +42,7 @@ pub(crate) struct MessageProcessor {
initialized: bool,
codex_linux_sandbox_exe: Option<PathBuf>,
conversation_manager: Arc<ConversationManager>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
}
impl MessageProcessor {
@@ -436,7 +437,10 @@ impl MessageProcessor {
tracing::info!("tools/call -> params: {:?}", arguments);
// parse arguments
let CodexToolCallReplyParam { session_id, prompt } = match arguments {
let CodexToolCallReplyParam {
conversation_id,
prompt,
} = match arguments {
Some(json_val) => match serde_json::from_value::<CodexToolCallReplyParam>(json_val) {
Ok(params) => params,
Err(e) => {
@@ -457,12 +461,12 @@ impl MessageProcessor {
},
None => {
tracing::error!(
"Missing arguments for codex-reply tool-call; the `session_id` and `prompt` fields are required."
"Missing arguments for codex-reply tool-call; the `conversation_id` and `prompt` fields are required."
);
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: "Missing arguments for codex-reply tool-call; the `session_id` and `prompt` fields are required.".to_owned(),
text: "Missing arguments for codex-reply tool-call; the `conversation_id` and `prompt` fields are required.".to_owned(),
annotations: None,
})],
is_error: Some(true),
@@ -473,14 +477,14 @@ impl MessageProcessor {
return;
}
};
let session_id = match Uuid::parse_str(&session_id) {
Ok(id) => id,
let conversation_id = match Uuid::parse_str(&conversation_id) {
Ok(id) => ConversationId::from(id),
Err(e) => {
tracing::error!("Failed to parse session_id: {e}");
tracing::error!("Failed to parse conversation_id: {e}");
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: format!("Failed to parse session_id: {e}"),
text: format!("Failed to parse conversation_id: {e}"),
annotations: None,
})],
is_error: Some(true),
@@ -496,14 +500,18 @@ impl MessageProcessor {
let outgoing = self.outgoing.clone();
let running_requests_id_to_codex_uuid = self.running_requests_id_to_codex_uuid.clone();
let codex = match self.conversation_manager.get_conversation(session_id).await {
let codex = match self
.conversation_manager
.get_conversation(conversation_id)
.await
{
Ok(c) => c,
Err(_) => {
tracing::warn!("Session not found for session_id: {session_id}");
tracing::warn!("Session not found for conversation_id: {conversation_id}");
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: format!("Session not found for session_id: {session_id}"),
text: format!("Session not found for conversation_id: {conversation_id}"),
annotations: None,
})],
is_error: Some(true),
@@ -528,7 +536,7 @@ impl MessageProcessor {
request_id,
prompt,
running_requests_id_to_codex_uuid,
session_id,
conversation_id,
)
.await;
}
@@ -564,24 +572,28 @@ impl MessageProcessor {
RequestId::Integer(i) => i.to_string(),
};
// Obtain the session_id while holding the first lock, then release.
let session_id = {
// Obtain the conversation id while holding the first lock, then release.
let conversation_id = {
let map_guard = self.running_requests_id_to_codex_uuid.lock().await;
match map_guard.get(&request_id) {
Some(id) => *id, // Uuid is Copy
Some(id) => *id,
None => {
tracing::warn!("Session not found for request_id: {}", request_id_string);
return;
}
}
};
tracing::info!("session_id: {session_id}");
tracing::info!("conversation_id: {conversation_id}");
// Obtain the Codex conversation from the server.
let codex_arc = match self.conversation_manager.get_conversation(session_id).await {
let codex_arc = match self
.conversation_manager
.get_conversation(conversation_id)
.await
{
Ok(c) => c,
Err(_) => {
tracing::warn!("Session not found for session_id: {session_id}");
tracing::warn!("Session not found for conversation_id: {conversation_id}");
return;
}
};

7
codex-rs/mcp-server/src/outgoing_message.rs
View File

@@ -258,6 +258,7 @@ pub(crate) struct OutgoingError {
mod tests {
use codex_core::protocol::EventMsg;
use codex_core::protocol::SessionConfiguredEvent;
use codex_protocol::mcp_protocol::ConversationId;
use codex_protocol::mcp_protocol::LoginChatGptCompleteNotification;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -270,10 +271,11 @@ mod tests {
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
let conversation_id = ConversationId::new();
let event = Event {
id: "1".to_string(),
msg: EventMsg::SessionConfigured(SessionConfiguredEvent {
session_id: Uuid::new_v4(),
session_id: conversation_id,
model: "gpt-4o".to_string(),
history_log_id: 1,
history_entry_count: 1000,
@@ -302,8 +304,9 @@ mod tests {
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
let conversation_id = ConversationId::new();
let session_configured_event = SessionConfiguredEvent {
session_id: Uuid::new_v4(),
session_id: conversation_id,
model: "gpt-4o".to_string(),
history_log_id: 1,
history_entry_count: 1000,