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

Add an elicitation for approve patch and refactor tool calls (#1642)

Browse Source 
1. Added an elicitation for `approve-patch` which is very similar to
`approve-exec`.
2. Extracted both elicitations to their own files to prevent
`codex_tool_runner` from blowing up in size.
This commit is contained in:
Gabriel Peal
2025-07-21 23:58:41 -07:00
committed by GitHub
parent 6cf4b96f9d
commit 710f728124
8 changed files with 584 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,38 +3,31 @@
//! and to make future feature-growth easier to manage.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use codex_core::Codex;
use codex_core::codex_wrapper::init_codex;
use codex_core::config::Config as CodexConfig;
use codex_core::protocol::AgentMessageEvent;
use codex_core::protocol::ApplyPatchApprovalRequestEvent;
use codex_core::protocol::EventMsg;
use codex_core::protocol::ExecApprovalRequestEvent;
use codex_core::protocol::InputItem;
use codex_core::protocol::Op;
use codex_core::protocol::ReviewDecision;
use codex_core::protocol::Submission;
use codex_core::protocol::TaskCompleteEvent;
use mcp_types::CallToolResult;
use mcp_types::ContentBlock;
use mcp_types::ElicitRequest;
use mcp_types::ElicitRequestParamsRequestedSchema;
use mcp_types::JSONRPCErrorError;
use mcp_types::ModelContextProtocolRequest;
use mcp_types::RequestId;
use mcp_types::TextContent;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use tokio::sync::Mutex;
use tracing::error;
use uuid::Uuid;
use crate::exec_approval::handle_exec_approval_request;
use crate::outgoing_message::OutgoingMessageSender;
use crate::patch_approval::handle_patch_approval_request;
const INVALID_PARAMS_ERROR_CODE: i64 = -32602;
pub(crate) const INVALID_PARAMS_ERROR_CODE: i64 = -32602;
/// Run a complete Codex session and stream events back to the client.
///
@@ -120,7 +113,7 @@ async fn run_codex_tool_session_inner(
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
) {
let sub_id = match &request_id {
let request_id_str = match &request_id {
RequestId::String(s) => s.clone(),
RequestId::Integer(n) => n.to_string(),
};
@@ -138,80 +131,34 @@ async fn run_codex_tool_session_inner(
cwd,
reason: _,
}) => {
let escaped_command = shlex::try_join(command.iter().map(|s| s.as_str()))
.unwrap_or_else(|_| command.join(" "));
let message = format!(
"Allow Codex to run `{escaped_command}` in `{cwd}`?",
cwd = cwd.to_string_lossy()
);
let params = ExecApprovalElicitRequestParams {
message,
requested_schema: ElicitRequestParamsRequestedSchema {
r#type: "object".to_string(),
properties: json!({}),
required: None,
},
codex_elicitation: "exec-approval".to_string(),
codex_mcp_tool_call_id: sub_id.clone(),
codex_event_id: event.id.clone(),
codex_command: command,
codex_cwd: cwd,
};
let params_json = match serde_json::to_value(&params) {
Ok(value) => value,
Err(err) => {
let message = format!(
"Failed to serialize ExecApprovalElicitRequestParams: {err}"
);
tracing::error!("{message}");
outgoing
.send_error(
request_id.clone(),
JSONRPCErrorError {
code: INVALID_PARAMS_ERROR_CODE,
message,
data: None,
},
)
.await;
continue;
}
};
let on_response = outgoing
.send_request(ElicitRequest::METHOD, Some(params_json))
.await;
// Listen for the response on a separate task so we do
// not block the main loop of this function.
{
let codex = codex.clone();
let event_id = event.id.clone();
tokio::spawn(async move {
on_exec_approval_response(event_id, on_response, codex).await;
});
}
// Continue, don't break so the session continues.
handle_exec_approval_request(
command,
cwd,
outgoing.clone(),
codex.clone(),
request_id.clone(),
request_id_str.clone(),
event.id.clone(),
)
.await;
continue;
}
EventMsg::ApplyPatchApprovalRequest(_) => {
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_string(),
text: "PATCH_APPROVAL_REQUIRED".to_string(),
annotations: None,
})],
is_error: None,
structured_content: None,
};
outgoing
.send_response(request_id.clone(), result.into())
.await;
// Continue, don't break so the session continues.
EventMsg::ApplyPatchApprovalRequest(ApplyPatchApprovalRequestEvent {
reason,
grant_root,
changes,
}) => {
handle_patch_approval_request(
reason,
grant_root,
changes,
outgoing.clone(),
codex.clone(),
request_id.clone(),
request_id_str.clone(),
event.id.clone(),
)
.await;
continue;
}
EventMsg::TaskComplete(TaskCompleteEvent { last_agent_message }) => {
@@ -286,71 +233,3 @@ async fn run_codex_tool_session_inner(
}
}
}
async fn on_exec_approval_response(
event_id: String,
receiver: tokio::sync::oneshot::Receiver<mcp_types::Result>,
codex: Arc<Codex>,
) {
let response = receiver.await;
let value = match response {
Ok(value) => value,
Err(err) => {
error!("request failed: {err:?}");
return;
}
};
// Try to deserialize `value` and then make the appropriate call to `codex`.
let response = match serde_json::from_value::<ExecApprovalResponse>(value) {
Ok(response) => response,
Err(err) => {
error!("failed to deserialize ExecApprovalResponse: {err}");
// If we cannot deserialize the response, we deny the request to be
// conservative.
ExecApprovalResponse {
decision: ReviewDecision::Denied,
}
}
};
if let Err(err) = codex
.submit(Op::ExecApproval {
id: event_id,
decision: response.decision,
})
.await
{
error!("failed to submit ExecApproval: {err}");
}
}
// TODO(mbolin): ExecApprovalResponse does not conform to ElicitResult. See:
// - https://github.com/modelcontextprotocol/modelcontextprotocol/blob/f962dc1780fa5eed7fb7c8a0232f1fc83ef220cd/schema/2025-06-18/schema.json#L617-L636
// - https://modelcontextprotocol.io/specification/draft/client/elicitation#protocol-messages
// It should have "action" and "content" fields.
#[derive(Debug, Serialize, Deserialize)]
pub struct ExecApprovalResponse {
pub decision: ReviewDecision,
}
/// Conforms to [`mcp_types::ElicitRequestParams`] so that it can be used as the
/// `params` field of an [`mcp_types::ElicitRequest`].
#[derive(Debug, Serialize)]
pub struct ExecApprovalElicitRequestParams {
// These fields are required so that `params`
// conforms to ElicitRequestParams.
pub message: String,
#[serde(rename = "requestedSchema")]
pub requested_schema: ElicitRequestParamsRequestedSchema,
// These are additional fields the client can use to
// correlate the request with the codex tool call.
pub codex_elicitation: String,
pub codex_mcp_tool_call_id: String,
pub codex_event_id: String,
pub codex_command: Vec<String>,
pub codex_cwd: PathBuf,
}

145
codex-rs/mcp-server/src/exec_approval.rs Normal file
View File

@@ -0,0 +1,145 @@
use std::path::PathBuf;
use std::sync::Arc;
use codex_core::Codex;
use codex_core::protocol::Op;
use codex_core::protocol::ReviewDecision;
use mcp_types::ElicitRequest;
use mcp_types::ElicitRequestParamsRequestedSchema;
use mcp_types::JSONRPCErrorError;
use mcp_types::ModelContextProtocolRequest;
use mcp_types::RequestId;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use tracing::error;
use crate::codex_tool_runner::INVALID_PARAMS_ERROR_CODE;
/// Conforms to [`mcp_types::ElicitRequestParams`] so that it can be used as the
/// `params` field of an [`ElicitRequest`].
#[derive(Debug, Serialize)]
pub struct ExecApprovalElicitRequestParams {
// These fields are required so that `params`
// conforms to ElicitRequestParams.
pub message: String,
#[serde(rename = "requestedSchema")]
pub requested_schema: ElicitRequestParamsRequestedSchema,
// These are additional fields the client can use to
// correlate the request with the codex tool call.
pub codex_elicitation: String,
pub codex_mcp_tool_call_id: String,
pub codex_event_id: String,
pub codex_command: Vec<String>,
pub codex_cwd: PathBuf,
}
// TODO(mbolin): ExecApprovalResponse does not conform to ElicitResult. See:
// - https://github.com/modelcontextprotocol/modelcontextprotocol/blob/f962dc1780fa5eed7fb7c8a0232f1fc83ef220cd/schema/2025-06-18/schema.json#L617-L636
// - https://modelcontextprotocol.io/specification/draft/client/elicitation#protocol-messages
// It should have "action" and "content" fields.
#[derive(Debug, Serialize, Deserialize)]
pub struct ExecApprovalResponse {
pub decision: ReviewDecision,
}
pub(crate) async fn handle_exec_approval_request(
command: Vec<String>,
cwd: PathBuf,
outgoing: Arc<crate::outgoing_message::OutgoingMessageSender>,
codex: Arc<Codex>,
request_id: RequestId,
tool_call_id: String,
event_id: String,
) {
let escaped_command =
shlex::try_join(command.iter().map(|s| s.as_str())).unwrap_or_else(|_| command.join(" "));
let message = format!(
"Allow Codex to run `{escaped_command}` in `{cwd}`?",
cwd = cwd.to_string_lossy()
);
let params = ExecApprovalElicitRequestParams {
message,
requested_schema: ElicitRequestParamsRequestedSchema {
r#type: "object".to_string(),
properties: json!({}),
required: None,
},
codex_elicitation: "exec-approval".to_string(),
codex_mcp_tool_call_id: tool_call_id.clone(),
codex_event_id: event_id.clone(),
codex_command: command,
codex_cwd: cwd,
};
let params_json = match serde_json::to_value(&params) {
Ok(value) => value,
Err(err) => {
let message = format!("Failed to serialize ExecApprovalElicitRequestParams: {err}");
error!("{message}");
outgoing
.send_error(
request_id.clone(),
JSONRPCErrorError {
code: INVALID_PARAMS_ERROR_CODE,
message,
data: None,
},
)
.await;
return;
}
};
let on_response = outgoing
.send_request(ElicitRequest::METHOD, Some(params_json))
.await;
// Listen for the response on a separate task so we don't block the main agent loop.
{
let codex = codex.clone();
let event_id = event_id.clone();
tokio::spawn(async move {
on_exec_approval_response(event_id, on_response, codex).await;
});
}
}
async fn on_exec_approval_response(
event_id: String,
receiver: tokio::sync::oneshot::Receiver<mcp_types::Result>,
codex: Arc<Codex>,
) {
let response = receiver.await;
let value = match response {
Ok(value) => value,
Err(err) => {
error!("request failed: {err:?}");
return;
}
};
// Try to deserialize `value` and then make the appropriate call to `codex`.
let response = serde_json::from_value::<ExecApprovalResponse>(value).unwrap_or_else(|err| {
error!("failed to deserialize ExecApprovalResponse: {err}");
// If we cannot deserialize the response, we deny the request to be
// conservative.
ExecApprovalResponse {
decision: ReviewDecision::Denied,
}
});
if let Err(err) = codex
.submit(Op::ExecApproval {
id: event_id,
decision: response.decision,
})
.await
{
error!("failed to submit ExecApproval: {err}");
}
}

8
codex-rs/mcp-server/src/lib.rs
View File

@@ -16,17 +16,21 @@ use tracing::info;
mod codex_tool_config;
mod codex_tool_runner;
mod exec_approval;
mod json_to_toml;
mod message_processor;
mod outgoing_message;
mod patch_approval;
use crate::message_processor::MessageProcessor;
use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::OutgoingMessageSender;
pub use crate::codex_tool_config::CodexToolCallParam;
pub use crate::codex_tool_runner::ExecApprovalElicitRequestParams;
pub use crate::codex_tool_runner::ExecApprovalResponse;
pub use crate::exec_approval::ExecApprovalElicitRequestParams;
pub use crate::exec_approval::ExecApprovalResponse;
pub use crate::patch_approval::PatchApprovalElicitRequestParams;
pub use crate::patch_approval::PatchApprovalResponse;
/// Size of the bounded channels used to communicate between tasks. The value
/// is a balance between throughput and memory usage – 128 messages should be

147
codex-rs/mcp-server/src/patch_approval.rs Normal file
View File

@@ -0,0 +1,147 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use codex_core::Codex;
use codex_core::protocol::FileChange;
use codex_core::protocol::Op;
use codex_core::protocol::ReviewDecision;
use mcp_types::ElicitRequest;
use mcp_types::ElicitRequestParamsRequestedSchema;
use mcp_types::JSONRPCErrorError;
use mcp_types::ModelContextProtocolRequest;
use mcp_types::RequestId;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use tracing::error;
use crate::codex_tool_runner::INVALID_PARAMS_ERROR_CODE;
use crate::outgoing_message::OutgoingMessageSender;
#[derive(Debug, Serialize)]
pub struct PatchApprovalElicitRequestParams {
pub message: String,
#[serde(rename = "requestedSchema")]
pub requested_schema: ElicitRequestParamsRequestedSchema,
pub codex_elicitation: String,
pub codex_mcp_tool_call_id: String,
pub codex_event_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub codex_reason: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub codex_grant_root: Option<PathBuf>,
pub codex_changes: HashMap<PathBuf, FileChange>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct PatchApprovalResponse {
pub decision: ReviewDecision,
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn handle_patch_approval_request(
reason: Option<String>,
grant_root: Option<PathBuf>,
changes: HashMap<PathBuf, FileChange>,
outgoing: Arc<OutgoingMessageSender>,
codex: Arc<Codex>,
request_id: RequestId,
tool_call_id: String,
event_id: String,
) {
let mut message_lines = Vec::new();
if let Some(r) = &reason {
message_lines.push(r.clone());
}
message_lines.push("Allow Codex to apply proposed code changes?".to_string());
let params = PatchApprovalElicitRequestParams {
message: message_lines.join("\n"),
requested_schema: ElicitRequestParamsRequestedSchema {
r#type: "object".to_string(),
properties: json!({}),
required: None,
},
codex_elicitation: "patch-approval".to_string(),
codex_mcp_tool_call_id: tool_call_id.clone(),
codex_event_id: event_id.clone(),
codex_reason: reason,
codex_grant_root: grant_root,
codex_changes: changes,
};
let params_json = match serde_json::to_value(&params) {
Ok(value) => value,
Err(err) => {
let message = format!("Failed to serialize PatchApprovalElicitRequestParams: {err}");
error!("{message}");
outgoing
.send_error(
request_id.clone(),
JSONRPCErrorError {
code: INVALID_PARAMS_ERROR_CODE,
message,
data: None,
},
)
.await;
return;
}
};
let on_response = outgoing
.send_request(ElicitRequest::METHOD, Some(params_json))
.await;
// Listen for the response on a separate task so we don't block the main agent loop.
{
let codex = codex.clone();
let event_id = event_id.clone();
tokio::spawn(async move {
on_patch_approval_response(event_id, on_response, codex).await;
});
}
}
pub(crate) async fn on_patch_approval_response(
event_id: String,
receiver: tokio::sync::oneshot::Receiver<mcp_types::Result>,
codex: Arc<Codex>,
) {
let response = receiver.await;
let value = match response {
Ok(value) => value,
Err(err) => {
error!("request failed: {err:?}");
if let Err(submit_err) = codex
.submit(Op::PatchApproval {
id: event_id.clone(),
decision: ReviewDecision::Denied,
})
.await
{
error!("failed to submit denied PatchApproval after request failure: {submit_err}");
}
return;
}
};
let response = serde_json::from_value::<PatchApprovalResponse>(value).unwrap_or_else(|err| {
error!("failed to deserialize PatchApprovalResponse: {err}");
PatchApprovalResponse {
decision: ReviewDecision::Denied,
}
});
if let Err(err) = codex
.submit(Op::PatchApproval {
id: event_id,
decision: response.decision,
})
.await
{
error!("failed to submit PatchApproval: {err}");
}
}