Phase 1: Repository & Infrastructure Setup
- Renamed directories: codex-rs -> llmx-rs, codex-cli -> llmx-cli
- Updated package.json files:
- Root: llmx-monorepo
- CLI: @llmx/llmx
- SDK: @llmx/llmx-sdk
- Updated pnpm workspace configuration
- Renamed binary: codex.js -> llmx.js
- Updated environment variables: CODEX_* -> LLMX_*
- Changed repository URLs to valknar/llmx
🤖 Generated with Claude Code
This commit is contained in:
Sebastian Krüger
46
llmx-rs/mcp-server/Cargo.toml
Normal file
46
llmx-rs/mcp-server/Cargo.toml
Normal file
@@ -0,0 +1,46 @@
|
||||
[package]
|
||||
edition = "2024"
|
||||
name = "codex-mcp-server"
|
||||
version = { workspace = true }
|
||||
|
||||
[[bin]]
|
||||
name = "codex-mcp-server"
|
||||
path = "src/main.rs"
|
||||
|
||||
[lib]
|
||||
name = "codex_mcp_server"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
codex-arg0 = { workspace = true }
|
||||
codex-common = { workspace = true, features = ["cli"] }
|
||||
codex-core = { workspace = true }
|
||||
codex-protocol = { workspace = true }
|
||||
codex-utils-json-to-toml = { workspace = true }
|
||||
mcp-types = { workspace = true }
|
||||
schemars = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
shlex = { workspace = true }
|
||||
tokio = { workspace = true, features = [
|
||||
"io-std",
|
||||
"macros",
|
||||
"process",
|
||||
"rt-multi-thread",
|
||||
"signal",
|
||||
] }
|
||||
tracing = { workspace = true, features = ["log"] }
|
||||
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
|
||||
|
||||
[dev-dependencies]
|
||||
assert_cmd = { workspace = true }
|
||||
core_test_support = { workspace = true }
|
||||
mcp_test_support = { workspace = true }
|
||||
os_info = { workspace = true }
|
||||
pretty_assertions = { workspace = true }
|
||||
tempfile = { workspace = true }
|
||||
wiremock = { workspace = true }
|
||||
347
llmx-rs/mcp-server/src/codex_tool_config.rs
Normal file
347
llmx-rs/mcp-server/src/codex_tool_config.rs
Normal file
@@ -0,0 +1,347 @@
|
||||
//! Configuration object accepted by the `codex` MCP tool-call.
|
||||
|
||||
use codex_core::protocol::AskForApproval;
|
||||
use codex_protocol::config_types::SandboxMode;
|
||||
use codex_utils_json_to_toml::json_to_toml;
|
||||
use mcp_types::Tool;
|
||||
use mcp_types::ToolInputSchema;
|
||||
use schemars::JsonSchema;
|
||||
use schemars::r#gen::SchemaSettings;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Client-supplied configuration for a `codex` tool-call.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub struct CodexToolCallParam {
|
||||
/// The *initial user prompt* to start the Codex conversation.
|
||||
pub prompt: String,
|
||||
|
||||
/// Optional override for the model name (e.g. "o3", "o4-mini").
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
|
||||
/// Configuration profile from config.toml to specify default options.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub profile: Option<String>,
|
||||
|
||||
/// Working directory for the session. If relative, it is resolved against
|
||||
/// the server process's current working directory.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cwd: Option<String>,
|
||||
|
||||
/// Approval policy for shell commands generated by the model:
|
||||
/// `untrusted`, `on-failure`, `on-request`, `never`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub approval_policy: Option<CodexToolCallApprovalPolicy>,
|
||||
|
||||
/// Sandbox mode: `read-only`, `workspace-write`, or `danger-full-access`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub sandbox: Option<CodexToolCallSandboxMode>,
|
||||
|
||||
/// Individual config settings that will override what is in
|
||||
/// CODEX_HOME/config.toml.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub config: Option<HashMap<String, serde_json::Value>>,
|
||||
|
||||
/// The set of instructions to use instead of the default ones.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub base_instructions: Option<String>,
|
||||
|
||||
/// Developer instructions that should be injected as a developer role message.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub developer_instructions: Option<String>,
|
||||
|
||||
/// Prompt used when compacting the conversation.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub compact_prompt: Option<String>,
|
||||
}
|
||||
|
||||
/// Custom enum mirroring [`AskForApproval`], but has an extra dependency on
|
||||
/// [`JsonSchema`].
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum CodexToolCallApprovalPolicy {
|
||||
Untrusted,
|
||||
OnFailure,
|
||||
OnRequest,
|
||||
Never,
|
||||
}
|
||||
|
||||
impl From<CodexToolCallApprovalPolicy> for AskForApproval {
|
||||
fn from(value: CodexToolCallApprovalPolicy) -> Self {
|
||||
match value {
|
||||
CodexToolCallApprovalPolicy::Untrusted => AskForApproval::UnlessTrusted,
|
||||
CodexToolCallApprovalPolicy::OnFailure => AskForApproval::OnFailure,
|
||||
CodexToolCallApprovalPolicy::OnRequest => AskForApproval::OnRequest,
|
||||
CodexToolCallApprovalPolicy::Never => AskForApproval::Never,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Custom enum mirroring [`SandboxMode`] from config_types.rs, but with
|
||||
/// `JsonSchema` support.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum CodexToolCallSandboxMode {
|
||||
ReadOnly,
|
||||
WorkspaceWrite,
|
||||
DangerFullAccess,
|
||||
}
|
||||
|
||||
impl From<CodexToolCallSandboxMode> for SandboxMode {
|
||||
fn from(value: CodexToolCallSandboxMode) -> Self {
|
||||
match value {
|
||||
CodexToolCallSandboxMode::ReadOnly => SandboxMode::ReadOnly,
|
||||
CodexToolCallSandboxMode::WorkspaceWrite => SandboxMode::WorkspaceWrite,
|
||||
CodexToolCallSandboxMode::DangerFullAccess => SandboxMode::DangerFullAccess,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a `Tool` definition (JSON schema etc.) for the Codex tool-call.
|
||||
pub(crate) fn create_tool_for_codex_tool_call_param() -> Tool {
|
||||
let schema = SchemaSettings::draft2019_09()
|
||||
.with(|s| {
|
||||
s.inline_subschemas = true;
|
||||
s.option_add_null_type = false;
|
||||
})
|
||||
.into_generator()
|
||||
.into_root_schema_for::<CodexToolCallParam>();
|
||||
|
||||
#[expect(clippy::expect_used)]
|
||||
let schema_value =
|
||||
serde_json::to_value(&schema).expect("Codex tool schema should serialise to JSON");
|
||||
|
||||
let tool_input_schema =
|
||||
serde_json::from_value::<ToolInputSchema>(schema_value).unwrap_or_else(|e| {
|
||||
panic!("failed to create Tool from schema: {e}");
|
||||
});
|
||||
|
||||
Tool {
|
||||
name: "codex".to_string(),
|
||||
title: Some("Codex".to_string()),
|
||||
input_schema: tool_input_schema,
|
||||
// TODO(mbolin): This should be defined.
|
||||
output_schema: None,
|
||||
description: Some(
|
||||
"Run a Codex session. Accepts configuration parameters matching the Codex Config struct.".to_string(),
|
||||
),
|
||||
annotations: None,
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexToolCallParam {
|
||||
/// Returns the initial user prompt to start the Codex conversation and the
|
||||
/// effective Config object generated from the supplied parameters.
|
||||
pub async fn into_config(
|
||||
self,
|
||||
codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
) -> std::io::Result<(String, codex_core::config::Config)> {
|
||||
let Self {
|
||||
prompt,
|
||||
model,
|
||||
profile,
|
||||
cwd,
|
||||
approval_policy,
|
||||
sandbox,
|
||||
config: cli_overrides,
|
||||
base_instructions,
|
||||
developer_instructions,
|
||||
compact_prompt,
|
||||
} = self;
|
||||
|
||||
// Build the `ConfigOverrides` recognized by codex-core.
|
||||
let overrides = codex_core::config::ConfigOverrides {
|
||||
model,
|
||||
review_model: None,
|
||||
config_profile: profile,
|
||||
cwd: cwd.map(PathBuf::from),
|
||||
approval_policy: approval_policy.map(Into::into),
|
||||
sandbox_mode: sandbox.map(Into::into),
|
||||
model_provider: None,
|
||||
codex_linux_sandbox_exe,
|
||||
base_instructions,
|
||||
developer_instructions,
|
||||
compact_prompt,
|
||||
include_apply_patch_tool: None,
|
||||
show_raw_agent_reasoning: None,
|
||||
tools_web_search_request: None,
|
||||
experimental_sandbox_command_assessment: None,
|
||||
additional_writable_roots: Vec::new(),
|
||||
};
|
||||
|
||||
let cli_overrides = cli_overrides
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(k, v)| (k, json_to_toml(v)))
|
||||
.collect();
|
||||
|
||||
let cfg =
|
||||
codex_core::config::Config::load_with_cli_overrides(cli_overrides, overrides).await?;
|
||||
|
||||
Ok((prompt, cfg))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CodexToolCallReplyParam {
|
||||
/// The conversation id for this Codex session.
|
||||
pub conversation_id: String,
|
||||
|
||||
/// The *next user prompt* to continue the Codex conversation.
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
/// Builds a `Tool` definition for the `codex-reply` tool-call.
|
||||
pub(crate) fn create_tool_for_codex_tool_call_reply_param() -> Tool {
|
||||
let schema = SchemaSettings::draft2019_09()
|
||||
.with(|s| {
|
||||
s.inline_subschemas = true;
|
||||
s.option_add_null_type = false;
|
||||
})
|
||||
.into_generator()
|
||||
.into_root_schema_for::<CodexToolCallReplyParam>();
|
||||
|
||||
#[expect(clippy::expect_used)]
|
||||
let schema_value =
|
||||
serde_json::to_value(&schema).expect("Codex reply tool schema should serialise to JSON");
|
||||
|
||||
let tool_input_schema =
|
||||
serde_json::from_value::<ToolInputSchema>(schema_value).unwrap_or_else(|e| {
|
||||
panic!("failed to create Tool from schema: {e}");
|
||||
});
|
||||
|
||||
Tool {
|
||||
name: "codex-reply".to_string(),
|
||||
title: Some("Codex Reply".to_string()),
|
||||
input_schema: tool_input_schema,
|
||||
output_schema: None,
|
||||
description: Some(
|
||||
"Continue a Codex conversation by providing the conversation id and prompt."
|
||||
.to_string(),
|
||||
),
|
||||
annotations: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
/// We include a test to verify the exact JSON schema as "executable
|
||||
/// documentation" for the schema. When can track changes to this test as a
|
||||
/// way to audit changes to the generated schema.
|
||||
///
|
||||
/// Seeing the fully expanded schema makes it easier to casually verify that
|
||||
/// the generated JSON for enum types such as "approval-policy" is compact.
|
||||
/// Ideally, modelcontextprotocol/inspector would provide a simpler UI for
|
||||
/// enum fields versus open string fields to take advantage of this.
|
||||
///
|
||||
/// As of 2025-05-04, there is an open PR for this:
|
||||
/// https://github.com/modelcontextprotocol/inspector/pull/196
|
||||
#[test]
|
||||
fn verify_codex_tool_json_schema() {
|
||||
let tool = create_tool_for_codex_tool_call_param();
|
||||
let tool_json = serde_json::to_value(&tool).expect("tool serializes");
|
||||
let expected_tool_json = serde_json::json!({
|
||||
"name": "codex",
|
||||
"title": "Codex",
|
||||
"description": "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"approval-policy": {
|
||||
"description": "Approval policy for shell commands generated by the model: `untrusted`, `on-failure`, `on-request`, `never`.",
|
||||
"enum": [
|
||||
"untrusted",
|
||||
"on-failure",
|
||||
"on-request",
|
||||
"never"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
"sandbox": {
|
||||
"description": "Sandbox mode: `read-only`, `workspace-write`, or `danger-full-access`.",
|
||||
"enum": [
|
||||
"read-only",
|
||||
"workspace-write",
|
||||
"danger-full-access"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
"config": {
|
||||
"description": "Individual config settings that will override what is in CODEX_HOME/config.toml.",
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
},
|
||||
"cwd": {
|
||||
"description": "Working directory for the session. If relative, it is resolved against the server process's current working directory.",
|
||||
"type": "string"
|
||||
},
|
||||
"model": {
|
||||
"description": "Optional override for the model name (e.g. \"o3\", \"o4-mini\").",
|
||||
"type": "string"
|
||||
},
|
||||
"profile": {
|
||||
"description": "Configuration profile from config.toml to specify default options.",
|
||||
"type": "string"
|
||||
},
|
||||
"prompt": {
|
||||
"description": "The *initial user prompt* to start the Codex conversation.",
|
||||
"type": "string"
|
||||
},
|
||||
"base-instructions": {
|
||||
"description": "The set of instructions to use instead of the default ones.",
|
||||
"type": "string"
|
||||
},
|
||||
"developer-instructions": {
|
||||
"description": "Developer instructions that should be injected as a developer role message.",
|
||||
"type": "string"
|
||||
},
|
||||
"compact-prompt": {
|
||||
"description": "Prompt used when compacting the conversation.",
|
||||
"type": "string"
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"prompt"
|
||||
]
|
||||
}
|
||||
});
|
||||
assert_eq!(expected_tool_json, tool_json);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_codex_tool_reply_json_schema() {
|
||||
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 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"
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"conversationId",
|
||||
"prompt",
|
||||
],
|
||||
"type": "object",
|
||||
},
|
||||
"name": "codex-reply",
|
||||
"title": "Codex Reply",
|
||||
});
|
||||
assert_eq!(expected_tool_json, tool_json);
|
||||
}
|
||||
}
|
||||
327
llmx-rs/mcp-server/src/codex_tool_runner.rs
Normal file
327
llmx-rs/mcp-server/src/codex_tool_runner.rs
Normal file
@@ -0,0 +1,327 @@
|
||||
//! Asynchronous worker that executes a **Codex** tool-call inside a spawned
|
||||
//! Tokio task. Separated from `message_processor.rs` to keep that file small
|
||||
//! and to make future feature-growth easier to manage.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::exec_approval::handle_exec_approval_request;
|
||||
use crate::outgoing_message::OutgoingMessageSender;
|
||||
use crate::outgoing_message::OutgoingNotificationMeta;
|
||||
use crate::patch_approval::handle_patch_approval_request;
|
||||
use codex_core::CodexConversation;
|
||||
use codex_core::ConversationManager;
|
||||
use codex_core::NewConversation;
|
||||
use codex_core::config::Config as CodexConfig;
|
||||
use codex_core::protocol::AgentMessageEvent;
|
||||
use codex_core::protocol::ApplyPatchApprovalRequestEvent;
|
||||
use codex_core::protocol::Event;
|
||||
use codex_core::protocol::EventMsg;
|
||||
use codex_core::protocol::ExecApprovalRequestEvent;
|
||||
use codex_core::protocol::Op;
|
||||
use codex_core::protocol::Submission;
|
||||
use codex_core::protocol::TaskCompleteEvent;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use mcp_types::CallToolResult;
|
||||
use mcp_types::ContentBlock;
|
||||
use mcp_types::RequestId;
|
||||
use mcp_types::TextContent;
|
||||
use serde_json::json;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
pub(crate) const INVALID_PARAMS_ERROR_CODE: i64 = -32602;
|
||||
|
||||
/// Run a complete Codex session and stream events back to the client.
|
||||
///
|
||||
/// On completion (success or error) the function sends the appropriate
|
||||
/// `tools/call` response so the LLM can continue the conversation.
|
||||
pub async fn run_codex_tool_session(
|
||||
id: RequestId,
|
||||
initial_prompt: String,
|
||||
config: CodexConfig,
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
conversation_manager: Arc<ConversationManager>,
|
||||
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
|
||||
) {
|
||||
let NewConversation {
|
||||
conversation_id,
|
||||
conversation,
|
||||
session_configured,
|
||||
} = match conversation_manager.new_conversation(config).await {
|
||||
Ok(res) => res,
|
||||
Err(e) => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_string(),
|
||||
text: format!("Failed to start Codex session: {e}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
outgoing.send_response(id.clone(), result).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let session_configured_event = Event {
|
||||
// Use a fake id value for now.
|
||||
id: "".to_string(),
|
||||
msg: EventMsg::SessionConfigured(session_configured.clone()),
|
||||
};
|
||||
outgoing
|
||||
.send_event_as_notification(
|
||||
&session_configured_event,
|
||||
Some(OutgoingNotificationMeta::new(Some(id.clone()))),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Use the original MCP request ID as the `sub_id` for the Codex submission so that
|
||||
// any events emitted for this tool-call can be correlated with the
|
||||
// originating `tools/call` request.
|
||||
let sub_id = match &id {
|
||||
RequestId::String(s) => s.clone(),
|
||||
RequestId::Integer(n) => n.to_string(),
|
||||
};
|
||||
running_requests_id_to_codex_uuid
|
||||
.lock()
|
||||
.await
|
||||
.insert(id.clone(), conversation_id);
|
||||
let submission = Submission {
|
||||
id: sub_id.clone(),
|
||||
op: Op::UserInput {
|
||||
items: vec![UserInput::Text {
|
||||
text: initial_prompt.clone(),
|
||||
}],
|
||||
},
|
||||
};
|
||||
|
||||
if let Err(e) = conversation.submit_with_id(submission).await {
|
||||
tracing::error!("Failed to submit initial prompt: {e}");
|
||||
// unregister the id so we don't keep it in the map
|
||||
running_requests_id_to_codex_uuid.lock().await.remove(&id);
|
||||
return;
|
||||
}
|
||||
|
||||
run_codex_tool_session_inner(
|
||||
conversation,
|
||||
outgoing,
|
||||
id,
|
||||
running_requests_id_to_codex_uuid,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub async fn run_codex_tool_session_reply(
|
||||
conversation: Arc<CodexConversation>,
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
request_id: RequestId,
|
||||
prompt: String,
|
||||
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(), conversation_id);
|
||||
if let Err(e) = conversation
|
||||
.submit(Op::UserInput {
|
||||
items: vec![UserInput::Text { text: prompt }],
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::error!("Failed to submit user input: {e}");
|
||||
// unregister the id so we don't keep it in the map
|
||||
running_requests_id_to_codex_uuid
|
||||
.lock()
|
||||
.await
|
||||
.remove(&request_id);
|
||||
return;
|
||||
}
|
||||
|
||||
run_codex_tool_session_inner(
|
||||
conversation,
|
||||
outgoing,
|
||||
request_id,
|
||||
running_requests_id_to_codex_uuid,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
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, ConversationId>>>,
|
||||
) {
|
||||
let request_id_str = match &request_id {
|
||||
RequestId::String(s) => s.clone(),
|
||||
RequestId::Integer(n) => n.to_string(),
|
||||
};
|
||||
|
||||
// Stream events until the task needs to pause for user interaction or
|
||||
// completes.
|
||||
loop {
|
||||
match codex.next_event().await {
|
||||
Ok(event) => {
|
||||
outgoing
|
||||
.send_event_as_notification(
|
||||
&event,
|
||||
Some(OutgoingNotificationMeta::new(Some(request_id.clone()))),
|
||||
)
|
||||
.await;
|
||||
|
||||
match event.msg {
|
||||
EventMsg::ExecApprovalRequest(ExecApprovalRequestEvent {
|
||||
command,
|
||||
cwd,
|
||||
call_id,
|
||||
reason: _,
|
||||
risk,
|
||||
parsed_cmd,
|
||||
}) => {
|
||||
handle_exec_approval_request(
|
||||
command,
|
||||
cwd,
|
||||
outgoing.clone(),
|
||||
codex.clone(),
|
||||
request_id.clone(),
|
||||
request_id_str.clone(),
|
||||
event.id.clone(),
|
||||
call_id,
|
||||
parsed_cmd,
|
||||
risk,
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
EventMsg::Error(err_event) => {
|
||||
// Return a response to conclude the tool call when the Codex session reports an error (e.g., interruption).
|
||||
let result = json!({
|
||||
"error": err_event.message,
|
||||
});
|
||||
outgoing.send_response(request_id.clone(), result).await;
|
||||
break;
|
||||
}
|
||||
EventMsg::Warning(_) => {
|
||||
continue;
|
||||
}
|
||||
EventMsg::ApplyPatchApprovalRequest(ApplyPatchApprovalRequestEvent {
|
||||
call_id,
|
||||
reason,
|
||||
grant_root,
|
||||
changes,
|
||||
}) => {
|
||||
handle_patch_approval_request(
|
||||
call_id,
|
||||
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 }) => {
|
||||
let text = match last_agent_message {
|
||||
Some(msg) => msg,
|
||||
None => "".to_string(),
|
||||
};
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_string(),
|
||||
text,
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: None,
|
||||
structured_content: None,
|
||||
};
|
||||
outgoing.send_response(request_id.clone(), result).await;
|
||||
// unregister the id so we don't keep it in the map
|
||||
running_requests_id_to_codex_uuid
|
||||
.lock()
|
||||
.await
|
||||
.remove(&request_id);
|
||||
break;
|
||||
}
|
||||
EventMsg::SessionConfigured(_) => {
|
||||
tracing::error!("unexpected SessionConfigured event");
|
||||
}
|
||||
EventMsg::AgentMessageDelta(_) => {
|
||||
// TODO: think how we want to support this in the MCP
|
||||
}
|
||||
EventMsg::AgentReasoningDelta(_) => {
|
||||
// TODO: think how we want to support this in the MCP
|
||||
}
|
||||
EventMsg::AgentMessage(AgentMessageEvent { .. }) => {
|
||||
// TODO: think how we want to support this in the MCP
|
||||
}
|
||||
EventMsg::AgentReasoningRawContent(_)
|
||||
| EventMsg::AgentReasoningRawContentDelta(_)
|
||||
| EventMsg::TaskStarted(_)
|
||||
| EventMsg::TokenCount(_)
|
||||
| EventMsg::AgentReasoning(_)
|
||||
| EventMsg::AgentReasoningSectionBreak(_)
|
||||
| EventMsg::McpToolCallBegin(_)
|
||||
| EventMsg::McpToolCallEnd(_)
|
||||
| EventMsg::McpListToolsResponse(_)
|
||||
| EventMsg::ListCustomPromptsResponse(_)
|
||||
| EventMsg::ExecCommandBegin(_)
|
||||
| EventMsg::ExecCommandOutputDelta(_)
|
||||
| EventMsg::ExecCommandEnd(_)
|
||||
| EventMsg::BackgroundEvent(_)
|
||||
| EventMsg::StreamError(_)
|
||||
| EventMsg::PatchApplyBegin(_)
|
||||
| EventMsg::PatchApplyEnd(_)
|
||||
| EventMsg::TurnDiff(_)
|
||||
| EventMsg::WebSearchBegin(_)
|
||||
| EventMsg::WebSearchEnd(_)
|
||||
| EventMsg::GetHistoryEntryResponse(_)
|
||||
| EventMsg::PlanUpdate(_)
|
||||
| EventMsg::TurnAborted(_)
|
||||
| EventMsg::UserMessage(_)
|
||||
| EventMsg::ShutdownComplete
|
||||
| EventMsg::ViewImageToolCall(_)
|
||||
| EventMsg::RawResponseItem(_)
|
||||
| EventMsg::EnteredReviewMode(_)
|
||||
| EventMsg::ItemStarted(_)
|
||||
| EventMsg::ItemCompleted(_)
|
||||
| EventMsg::AgentMessageContentDelta(_)
|
||||
| EventMsg::ReasoningContentDelta(_)
|
||||
| EventMsg::ReasoningRawContentDelta(_)
|
||||
| EventMsg::UndoStarted(_)
|
||||
| EventMsg::UndoCompleted(_)
|
||||
| EventMsg::ExitedReviewMode(_)
|
||||
| EventMsg::DeprecationNotice(_) => {
|
||||
// For now, we do not do anything extra for these
|
||||
// events. Note that
|
||||
// send(codex_event_to_notification(&event)) above has
|
||||
// already dispatched these events as notifications,
|
||||
// though we may want to do give different treatment to
|
||||
// individual events in the future.
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_string(),
|
||||
text: format!("Codex runtime error: {e}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
// TODO(mbolin): Could present the error in a more
|
||||
// structured way.
|
||||
structured_content: None,
|
||||
};
|
||||
outgoing.send_response(request_id.clone(), result).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
2
llmx-rs/mcp-server/src/error_code.rs
Normal file
2
llmx-rs/mcp-server/src/error_code.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
pub(crate) const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
|
||||
pub(crate) const INTERNAL_ERROR_CODE: i64 = -32603;
|
||||
158
llmx-rs/mcp-server/src/exec_approval.rs
Normal file
158
llmx-rs/mcp-server/src/exec_approval.rs
Normal file
@@ -0,0 +1,158 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_core::CodexConversation;
|
||||
use codex_core::protocol::Op;
|
||||
use codex_core::protocol::ReviewDecision;
|
||||
use codex_core::protocol::SandboxCommandAssessment;
|
||||
use codex_protocol::parse_command::ParsedCommand;
|
||||
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, Deserialize, 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_call_id: String,
|
||||
pub codex_command: Vec<String>,
|
||||
pub codex_cwd: PathBuf,
|
||||
pub codex_parsed_cmd: Vec<ParsedCommand>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub codex_risk: Option<SandboxCommandAssessment>,
|
||||
}
|
||||
|
||||
// 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,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn handle_exec_approval_request(
|
||||
command: Vec<String>,
|
||||
cwd: PathBuf,
|
||||
outgoing: Arc<crate::outgoing_message::OutgoingMessageSender>,
|
||||
codex: Arc<CodexConversation>,
|
||||
request_id: RequestId,
|
||||
tool_call_id: String,
|
||||
event_id: String,
|
||||
call_id: String,
|
||||
codex_parsed_cmd: Vec<ParsedCommand>,
|
||||
codex_risk: Option<SandboxCommandAssessment>,
|
||||
) {
|
||||
let escaped_command =
|
||||
shlex::try_join(command.iter().map(String::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_call_id: call_id,
|
||||
codex_command: command,
|
||||
codex_cwd: cwd,
|
||||
codex_parsed_cmd,
|
||||
codex_risk,
|
||||
};
|
||||
let params_json = match serde_json::to_value(¶ms) {
|
||||
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<CodexConversation>,
|
||||
) {
|
||||
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}");
|
||||
}
|
||||
}
|
||||
150
llmx-rs/mcp-server/src/lib.rs
Normal file
150
llmx-rs/mcp-server/src/lib.rs
Normal file
@@ -0,0 +1,150 @@
|
||||
//! Prototype MCP server.
|
||||
#![deny(clippy::print_stdout, clippy::print_stderr)]
|
||||
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Result as IoResult;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use codex_common::CliConfigOverrides;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigOverrides;
|
||||
|
||||
use mcp_types::JSONRPCMessage;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::io::{self};
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::debug;
|
||||
use tracing::error;
|
||||
use tracing::info;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
mod codex_tool_config;
|
||||
mod codex_tool_runner;
|
||||
mod error_code;
|
||||
mod exec_approval;
|
||||
pub(crate) 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_config::CodexToolCallReplyParam;
|
||||
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
|
||||
/// plenty for an interactive CLI.
|
||||
const CHANNEL_CAPACITY: usize = 128;
|
||||
|
||||
pub async fn run_main(
|
||||
codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
cli_config_overrides: CliConfigOverrides,
|
||||
) -> IoResult<()> {
|
||||
// Install a simple subscriber so `tracing` output is visible. Users can
|
||||
// control the log level with `RUST_LOG`.
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_env_filter(EnvFilter::from_default_env())
|
||||
.init();
|
||||
|
||||
// Set up channels.
|
||||
let (incoming_tx, mut incoming_rx) = mpsc::channel::<JSONRPCMessage>(CHANNEL_CAPACITY);
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
|
||||
|
||||
// Task: read from stdin, push to `incoming_tx`.
|
||||
let stdin_reader_handle = tokio::spawn({
|
||||
async move {
|
||||
let stdin = io::stdin();
|
||||
let reader = BufReader::new(stdin);
|
||||
let mut lines = reader.lines();
|
||||
|
||||
while let Some(line) = lines.next_line().await.unwrap_or_default() {
|
||||
match serde_json::from_str::<JSONRPCMessage>(&line) {
|
||||
Ok(msg) => {
|
||||
if incoming_tx.send(msg).await.is_err() {
|
||||
// Receiver gone – nothing left to do.
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => error!("Failed to deserialize JSONRPCMessage: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
debug!("stdin reader finished (EOF)");
|
||||
}
|
||||
});
|
||||
|
||||
// Parse CLI overrides once and derive the base Config eagerly so later
|
||||
// components do not need to work with raw TOML values.
|
||||
let cli_kv_overrides = cli_config_overrides.parse_overrides().map_err(|e| {
|
||||
std::io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!("error parsing -c overrides: {e}"),
|
||||
)
|
||||
})?;
|
||||
let config = Config::load_with_cli_overrides(cli_kv_overrides, ConfigOverrides::default())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}"))
|
||||
})?;
|
||||
|
||||
// Task: process incoming messages.
|
||||
let processor_handle = tokio::spawn({
|
||||
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
|
||||
let mut processor = MessageProcessor::new(
|
||||
outgoing_message_sender,
|
||||
codex_linux_sandbox_exe,
|
||||
std::sync::Arc::new(config),
|
||||
);
|
||||
async move {
|
||||
while let Some(msg) = incoming_rx.recv().await {
|
||||
match msg {
|
||||
JSONRPCMessage::Request(r) => processor.process_request(r).await,
|
||||
JSONRPCMessage::Response(r) => processor.process_response(r).await,
|
||||
JSONRPCMessage::Notification(n) => processor.process_notification(n).await,
|
||||
JSONRPCMessage::Error(e) => processor.process_error(e),
|
||||
}
|
||||
}
|
||||
|
||||
info!("processor task exited (channel closed)");
|
||||
}
|
||||
});
|
||||
|
||||
// Task: write outgoing messages to stdout.
|
||||
let stdout_writer_handle = tokio::spawn(async move {
|
||||
let mut stdout = io::stdout();
|
||||
while let Some(outgoing_message) = outgoing_rx.recv().await {
|
||||
let msg: JSONRPCMessage = outgoing_message.into();
|
||||
match serde_json::to_string(&msg) {
|
||||
Ok(json) => {
|
||||
if let Err(e) = stdout.write_all(json.as_bytes()).await {
|
||||
error!("Failed to write to stdout: {e}");
|
||||
break;
|
||||
}
|
||||
if let Err(e) = stdout.write_all(b"\n").await {
|
||||
error!("Failed to write newline to stdout: {e}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => error!("Failed to serialize JSONRPCMessage: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
info!("stdout writer exited (channel closed)");
|
||||
});
|
||||
|
||||
// Wait for all tasks to finish. The typical exit path is the stdin reader
|
||||
// hitting EOF which, once it drops `incoming_tx`, propagates shutdown to
|
||||
// the processor and then to the stdout task.
|
||||
let _ = tokio::join!(stdin_reader_handle, processor_handle, stdout_writer_handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
10
llmx-rs/mcp-server/src/main.rs
Normal file
10
llmx-rs/mcp-server/src/main.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
use codex_arg0::arg0_dispatch_or_else;
|
||||
use codex_common::CliConfigOverrides;
|
||||
use codex_mcp_server::run_main;
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
|
||||
run_main(codex_linux_sandbox_exe, CliConfigOverrides::default()).await?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
655
llmx-rs/mcp-server/src/message_processor.rs
Normal file
655
llmx-rs/mcp-server/src/message_processor.rs
Normal file
@@ -0,0 +1,655 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::codex_tool_config::CodexToolCallParam;
|
||||
use crate::codex_tool_config::CodexToolCallReplyParam;
|
||||
use crate::codex_tool_config::create_tool_for_codex_tool_call_param;
|
||||
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::ConversationId;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::ConversationManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::default_client::USER_AGENT_SUFFIX;
|
||||
use codex_core::default_client::get_codex_user_agent;
|
||||
use codex_core::protocol::Submission;
|
||||
use mcp_types::CallToolRequestParams;
|
||||
use mcp_types::CallToolResult;
|
||||
use mcp_types::ClientRequest as McpClientRequest;
|
||||
use mcp_types::ContentBlock;
|
||||
use mcp_types::JSONRPCError;
|
||||
use mcp_types::JSONRPCErrorError;
|
||||
use mcp_types::JSONRPCNotification;
|
||||
use mcp_types::JSONRPCRequest;
|
||||
use mcp_types::JSONRPCResponse;
|
||||
use mcp_types::ListToolsResult;
|
||||
use mcp_types::ModelContextProtocolRequest;
|
||||
use mcp_types::RequestId;
|
||||
use mcp_types::ServerCapabilitiesTools;
|
||||
use mcp_types::ServerNotification;
|
||||
use mcp_types::TextContent;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::task;
|
||||
|
||||
pub(crate) struct MessageProcessor {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
initialized: bool,
|
||||
codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
conversation_manager: Arc<ConversationManager>,
|
||||
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ConversationId>>>,
|
||||
}
|
||||
|
||||
impl MessageProcessor {
|
||||
/// Create a new `MessageProcessor`, retaining a handle to the outgoing
|
||||
/// `Sender` so handlers can enqueue messages to be written to stdout.
|
||||
pub(crate) fn new(
|
||||
outgoing: OutgoingMessageSender,
|
||||
codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
config: Arc<Config>,
|
||||
) -> Self {
|
||||
let outgoing = Arc::new(outgoing);
|
||||
let auth_manager = AuthManager::shared(
|
||||
config.codex_home.clone(),
|
||||
false,
|
||||
config.cli_auth_credentials_store_mode,
|
||||
);
|
||||
let conversation_manager =
|
||||
Arc::new(ConversationManager::new(auth_manager, SessionSource::Mcp));
|
||||
Self {
|
||||
outgoing,
|
||||
initialized: false,
|
||||
codex_linux_sandbox_exe,
|
||||
conversation_manager,
|
||||
running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn process_request(&mut self, request: JSONRPCRequest) {
|
||||
// Hold on to the ID so we can respond.
|
||||
let request_id = request.id.clone();
|
||||
|
||||
let client_request = match McpClientRequest::try_from(request) {
|
||||
Ok(client_request) => client_request,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to convert request: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Dispatch to a dedicated handler for each request type.
|
||||
match client_request {
|
||||
McpClientRequest::InitializeRequest(params) => {
|
||||
self.handle_initialize(request_id, params).await;
|
||||
}
|
||||
McpClientRequest::PingRequest(params) => {
|
||||
self.handle_ping(request_id, params).await;
|
||||
}
|
||||
McpClientRequest::ListResourcesRequest(params) => {
|
||||
self.handle_list_resources(params);
|
||||
}
|
||||
McpClientRequest::ListResourceTemplatesRequest(params) => {
|
||||
self.handle_list_resource_templates(params);
|
||||
}
|
||||
McpClientRequest::ReadResourceRequest(params) => {
|
||||
self.handle_read_resource(params);
|
||||
}
|
||||
McpClientRequest::SubscribeRequest(params) => {
|
||||
self.handle_subscribe(params);
|
||||
}
|
||||
McpClientRequest::UnsubscribeRequest(params) => {
|
||||
self.handle_unsubscribe(params);
|
||||
}
|
||||
McpClientRequest::ListPromptsRequest(params) => {
|
||||
self.handle_list_prompts(params);
|
||||
}
|
||||
McpClientRequest::GetPromptRequest(params) => {
|
||||
self.handle_get_prompt(params);
|
||||
}
|
||||
McpClientRequest::ListToolsRequest(params) => {
|
||||
self.handle_list_tools(request_id, params).await;
|
||||
}
|
||||
McpClientRequest::CallToolRequest(params) => {
|
||||
self.handle_call_tool(request_id, params).await;
|
||||
}
|
||||
McpClientRequest::SetLevelRequest(params) => {
|
||||
self.handle_set_level(params);
|
||||
}
|
||||
McpClientRequest::CompleteRequest(params) => {
|
||||
self.handle_complete(params);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a standalone JSON-RPC response originating from the peer.
|
||||
pub(crate) async fn process_response(&mut self, response: JSONRPCResponse) {
|
||||
tracing::info!("<- response: {:?}", response);
|
||||
let JSONRPCResponse { id, result, .. } = response;
|
||||
self.outgoing.notify_client_response(id, result).await
|
||||
}
|
||||
|
||||
/// Handle a fire-and-forget JSON-RPC notification.
|
||||
pub(crate) async fn process_notification(&mut self, notification: JSONRPCNotification) {
|
||||
let server_notification = match ServerNotification::try_from(notification) {
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to convert notification: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Similar to requests, route each notification type to its own stub
|
||||
// handler so additional logic can be implemented incrementally.
|
||||
match server_notification {
|
||||
ServerNotification::CancelledNotification(params) => {
|
||||
self.handle_cancelled_notification(params).await;
|
||||
}
|
||||
ServerNotification::ProgressNotification(params) => {
|
||||
self.handle_progress_notification(params);
|
||||
}
|
||||
ServerNotification::ResourceListChangedNotification(params) => {
|
||||
self.handle_resource_list_changed(params);
|
||||
}
|
||||
ServerNotification::ResourceUpdatedNotification(params) => {
|
||||
self.handle_resource_updated(params);
|
||||
}
|
||||
ServerNotification::PromptListChangedNotification(params) => {
|
||||
self.handle_prompt_list_changed(params);
|
||||
}
|
||||
ServerNotification::ToolListChangedNotification(params) => {
|
||||
self.handle_tool_list_changed(params);
|
||||
}
|
||||
ServerNotification::LoggingMessageNotification(params) => {
|
||||
self.handle_logging_message(params);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle an error object received from the peer.
|
||||
pub(crate) fn process_error(&mut self, err: JSONRPCError) {
|
||||
tracing::error!("<- error: {:?}", err);
|
||||
}
|
||||
|
||||
async fn handle_initialize(
|
||||
&mut self,
|
||||
id: RequestId,
|
||||
params: <mcp_types::InitializeRequest as ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("initialize -> params: {:?}", params);
|
||||
|
||||
if self.initialized {
|
||||
// Already initialised: send JSON-RPC error response.
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "initialize called more than once".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(id, error).await;
|
||||
return;
|
||||
}
|
||||
|
||||
let client_info = params.client_info;
|
||||
let name = client_info.name;
|
||||
let version = client_info.version;
|
||||
let user_agent_suffix = format!("{name}; {version}");
|
||||
if let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() {
|
||||
*suffix = Some(user_agent_suffix);
|
||||
}
|
||||
|
||||
self.initialized = true;
|
||||
|
||||
// Build a minimal InitializeResult. Fill with placeholders.
|
||||
let result = mcp_types::InitializeResult {
|
||||
capabilities: mcp_types::ServerCapabilities {
|
||||
completions: None,
|
||||
experimental: None,
|
||||
logging: None,
|
||||
prompts: None,
|
||||
resources: None,
|
||||
tools: Some(ServerCapabilitiesTools {
|
||||
list_changed: Some(true),
|
||||
}),
|
||||
},
|
||||
instructions: None,
|
||||
protocol_version: params.protocol_version.clone(),
|
||||
server_info: mcp_types::Implementation {
|
||||
name: "codex-mcp-server".to_string(),
|
||||
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
title: Some("Codex".to_string()),
|
||||
user_agent: Some(get_codex_user_agent()),
|
||||
},
|
||||
};
|
||||
|
||||
self.send_response::<mcp_types::InitializeRequest>(id, result)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn send_response<T>(&self, id: RequestId, result: T::Result)
|
||||
where
|
||||
T: ModelContextProtocolRequest,
|
||||
{
|
||||
self.outgoing.send_response(id, result).await;
|
||||
}
|
||||
|
||||
async fn handle_ping(
|
||||
&self,
|
||||
id: RequestId,
|
||||
params: <mcp_types::PingRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("ping -> params: {:?}", params);
|
||||
let result = json!({});
|
||||
self.send_response::<mcp_types::PingRequest>(id, result)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn handle_list_resources(
|
||||
&self,
|
||||
params: <mcp_types::ListResourcesRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("resources/list -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_list_resource_templates(
|
||||
&self,
|
||||
params:
|
||||
<mcp_types::ListResourceTemplatesRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("resources/templates/list -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_read_resource(
|
||||
&self,
|
||||
params: <mcp_types::ReadResourceRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("resources/read -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_subscribe(
|
||||
&self,
|
||||
params: <mcp_types::SubscribeRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("resources/subscribe -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_unsubscribe(
|
||||
&self,
|
||||
params: <mcp_types::UnsubscribeRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("resources/unsubscribe -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_list_prompts(
|
||||
&self,
|
||||
params: <mcp_types::ListPromptsRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("prompts/list -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_get_prompt(
|
||||
&self,
|
||||
params: <mcp_types::GetPromptRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("prompts/get -> params: {:?}", params);
|
||||
}
|
||||
|
||||
async fn handle_list_tools(
|
||||
&self,
|
||||
id: RequestId,
|
||||
params: <mcp_types::ListToolsRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::trace!("tools/list -> {params:?}");
|
||||
let result = ListToolsResult {
|
||||
tools: vec![
|
||||
create_tool_for_codex_tool_call_param(),
|
||||
create_tool_for_codex_tool_call_reply_param(),
|
||||
],
|
||||
next_cursor: None,
|
||||
};
|
||||
|
||||
self.send_response::<mcp_types::ListToolsRequest>(id, result)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn handle_call_tool(
|
||||
&self,
|
||||
id: RequestId,
|
||||
params: <mcp_types::CallToolRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("tools/call -> params: {:?}", params);
|
||||
let CallToolRequestParams { name, arguments } = params;
|
||||
|
||||
match name.as_str() {
|
||||
"codex" => self.handle_tool_call_codex(id, arguments).await,
|
||||
"codex-reply" => {
|
||||
self.handle_tool_call_codex_session_reply(id, arguments)
|
||||
.await
|
||||
}
|
||||
_ => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_string(),
|
||||
text: format!("Unknown tool '{name}'"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(id, result)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn handle_tool_call_codex(&self, id: RequestId, arguments: Option<serde_json::Value>) {
|
||||
let (initial_prompt, config): (String, Config) = match arguments {
|
||||
Some(json_val) => match serde_json::from_value::<CodexToolCallParam>(json_val) {
|
||||
Ok(tool_cfg) => match tool_cfg
|
||||
.into_config(self.codex_linux_sandbox_exe.clone())
|
||||
.await
|
||||
{
|
||||
Ok(cfg) => cfg,
|
||||
Err(e) => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_owned(),
|
||||
text: format!(
|
||||
"Failed to load Codex configuration from overrides: {e}"
|
||||
),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_owned(),
|
||||
text: format!("Failed to parse configuration for Codex tool: {e}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
},
|
||||
None => {
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_string(),
|
||||
text:
|
||||
"Missing arguments for codex tool-call; the `prompt` field is required."
|
||||
.to_string(),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Clone outgoing and server to move into async task.
|
||||
let outgoing = self.outgoing.clone();
|
||||
let conversation_manager = self.conversation_manager.clone();
|
||||
let running_requests_id_to_codex_uuid = self.running_requests_id_to_codex_uuid.clone();
|
||||
|
||||
// Spawn an async task to handle the Codex session so that we do not
|
||||
// block the synchronous message-processing loop.
|
||||
task::spawn(async move {
|
||||
// Run the Codex session and stream events back to the client.
|
||||
crate::codex_tool_runner::run_codex_tool_session(
|
||||
id,
|
||||
initial_prompt,
|
||||
config,
|
||||
outgoing,
|
||||
conversation_manager,
|
||||
running_requests_id_to_codex_uuid,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
async fn handle_tool_call_codex_session_reply(
|
||||
&self,
|
||||
request_id: RequestId,
|
||||
arguments: Option<serde_json::Value>,
|
||||
) {
|
||||
tracing::info!("tools/call -> params: {:?}", arguments);
|
||||
|
||||
// parse arguments
|
||||
let CodexToolCallReplyParam {
|
||||
conversation_id,
|
||||
prompt,
|
||||
} = match arguments {
|
||||
Some(json_val) => match serde_json::from_value::<CodexToolCallReplyParam>(json_val) {
|
||||
Ok(params) => params,
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to parse Codex tool call reply parameters: {e}");
|
||||
let result = CallToolResult {
|
||||
content: vec![ContentBlock::TextContent(TextContent {
|
||||
r#type: "text".to_owned(),
|
||||
text: format!("Failed to parse configuration for Codex tool: {e}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
},
|
||||
None => {
|
||||
tracing::error!(
|
||||
"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 `conversation_id` and `prompt` fields are required.".to_owned(),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let conversation_id = match ConversationId::from_string(&conversation_id) {
|
||||
Ok(id) => id,
|
||||
Err(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 conversation_id: {e}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Clone outgoing to move into async task.
|
||||
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(conversation_id)
|
||||
.await
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(_) => {
|
||||
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 conversation_id: {conversation_id}"),
|
||||
annotations: None,
|
||||
})],
|
||||
is_error: Some(true),
|
||||
structured_content: None,
|
||||
};
|
||||
outgoing.send_response(request_id, result).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Spawn the long-running reply handler.
|
||||
tokio::spawn({
|
||||
let outgoing = outgoing.clone();
|
||||
let prompt = prompt.clone();
|
||||
let running_requests_id_to_codex_uuid = running_requests_id_to_codex_uuid.clone();
|
||||
|
||||
async move {
|
||||
crate::codex_tool_runner::run_codex_tool_session_reply(
|
||||
codex,
|
||||
outgoing,
|
||||
request_id,
|
||||
prompt,
|
||||
running_requests_id_to_codex_uuid,
|
||||
conversation_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn handle_set_level(
|
||||
&self,
|
||||
params: <mcp_types::SetLevelRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("logging/setLevel -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_complete(
|
||||
&self,
|
||||
params: <mcp_types::CompleteRequest as mcp_types::ModelContextProtocolRequest>::Params,
|
||||
) {
|
||||
tracing::info!("completion/complete -> params: {:?}", params);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// Notification handlers
|
||||
// ---------------------------------------------------------------------
|
||||
|
||||
async fn handle_cancelled_notification(
|
||||
&self,
|
||||
params: <mcp_types::CancelledNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
let request_id = params.request_id;
|
||||
// Create a stable string form early for logging and submission id.
|
||||
let request_id_string = match &request_id {
|
||||
RequestId::String(s) => s.clone(),
|
||||
RequestId::Integer(i) => i.to_string(),
|
||||
};
|
||||
|
||||
// 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,
|
||||
None => {
|
||||
tracing::warn!("Session not found for request_id: {}", request_id_string);
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
tracing::info!("conversation_id: {conversation_id}");
|
||||
|
||||
// Obtain the Codex conversation from the server.
|
||||
let codex_arc = match self
|
||||
.conversation_manager
|
||||
.get_conversation(conversation_id)
|
||||
.await
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(_) => {
|
||||
tracing::warn!("Session not found for conversation_id: {conversation_id}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Submit interrupt to Codex.
|
||||
let err = codex_arc
|
||||
.submit_with_id(Submission {
|
||||
id: request_id_string,
|
||||
op: codex_core::protocol::Op::Interrupt,
|
||||
})
|
||||
.await;
|
||||
if let Err(e) = err {
|
||||
tracing::error!("Failed to submit interrupt to Codex: {e}");
|
||||
return;
|
||||
}
|
||||
// unregister the id so we don't keep it in the map
|
||||
self.running_requests_id_to_codex_uuid
|
||||
.lock()
|
||||
.await
|
||||
.remove(&request_id);
|
||||
}
|
||||
|
||||
fn handle_progress_notification(
|
||||
&self,
|
||||
params: <mcp_types::ProgressNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!("notifications/progress -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_resource_list_changed(
|
||||
&self,
|
||||
params: <mcp_types::ResourceListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!(
|
||||
"notifications/resources/list_changed -> params: {:?}",
|
||||
params
|
||||
);
|
||||
}
|
||||
|
||||
fn handle_resource_updated(
|
||||
&self,
|
||||
params: <mcp_types::ResourceUpdatedNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!("notifications/resources/updated -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_prompt_list_changed(
|
||||
&self,
|
||||
params: <mcp_types::PromptListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!("notifications/prompts/list_changed -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_tool_list_changed(
|
||||
&self,
|
||||
params: <mcp_types::ToolListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!("notifications/tools/list_changed -> params: {:?}", params);
|
||||
}
|
||||
|
||||
fn handle_logging_message(
|
||||
&self,
|
||||
params: <mcp_types::LoggingMessageNotification as mcp_types::ModelContextProtocolNotification>::Params,
|
||||
) {
|
||||
tracing::info!("notifications/message -> params: {:?}", params);
|
||||
}
|
||||
}
|
||||
333
llmx-rs/mcp-server/src/outgoing_message.rs
Normal file
333
llmx-rs/mcp-server/src/outgoing_message.rs
Normal file
@@ -0,0 +1,333 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::AtomicI64;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use codex_core::protocol::Event;
|
||||
use mcp_types::JSONRPC_VERSION;
|
||||
use mcp_types::JSONRPCError;
|
||||
use mcp_types::JSONRPCErrorError;
|
||||
use mcp_types::JSONRPCMessage;
|
||||
use mcp_types::JSONRPCNotification;
|
||||
use mcp_types::JSONRPCRequest;
|
||||
use mcp_types::JSONRPCResponse;
|
||||
use mcp_types::RequestId;
|
||||
use mcp_types::Result;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::error_code::INTERNAL_ERROR_CODE;
|
||||
|
||||
/// Sends messages to the client and manages request callbacks.
|
||||
pub(crate) struct OutgoingMessageSender {
|
||||
next_request_id: AtomicI64,
|
||||
sender: mpsc::UnboundedSender<OutgoingMessage>,
|
||||
request_id_to_callback: Mutex<HashMap<RequestId, oneshot::Sender<Result>>>,
|
||||
}
|
||||
|
||||
impl OutgoingMessageSender {
|
||||
pub(crate) fn new(sender: mpsc::UnboundedSender<OutgoingMessage>) -> Self {
|
||||
Self {
|
||||
next_request_id: AtomicI64::new(0),
|
||||
sender,
|
||||
request_id_to_callback: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn send_request(
|
||||
&self,
|
||||
method: &str,
|
||||
params: Option<serde_json::Value>,
|
||||
) -> oneshot::Receiver<Result> {
|
||||
let id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::Relaxed));
|
||||
let outgoing_message_id = id.clone();
|
||||
let (tx_approve, rx_approve) = oneshot::channel();
|
||||
{
|
||||
let mut request_id_to_callback = self.request_id_to_callback.lock().await;
|
||||
request_id_to_callback.insert(id, tx_approve);
|
||||
}
|
||||
|
||||
let outgoing_message = OutgoingMessage::Request(OutgoingRequest {
|
||||
id: outgoing_message_id,
|
||||
method: method.to_string(),
|
||||
params,
|
||||
});
|
||||
let _ = self.sender.send(outgoing_message);
|
||||
rx_approve
|
||||
}
|
||||
|
||||
pub(crate) async fn notify_client_response(&self, id: RequestId, result: Result) {
|
||||
let entry = {
|
||||
let mut request_id_to_callback = self.request_id_to_callback.lock().await;
|
||||
request_id_to_callback.remove_entry(&id)
|
||||
};
|
||||
|
||||
match entry {
|
||||
Some((id, sender)) => {
|
||||
if let Err(err) = sender.send(result) {
|
||||
warn!("could not notify callback for {id:?} due to: {err:?}");
|
||||
}
|
||||
}
|
||||
None => {
|
||||
warn!("could not find callback for {id:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn send_response<T: Serialize>(&self, id: RequestId, response: T) {
|
||||
match serde_json::to_value(response) {
|
||||
Ok(result) => {
|
||||
let outgoing_message = OutgoingMessage::Response(OutgoingResponse { id, result });
|
||||
let _ = self.sender.send(outgoing_message);
|
||||
}
|
||||
Err(err) => {
|
||||
self.send_error(
|
||||
id,
|
||||
JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to serialize response: {err}"),
|
||||
data: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// This is used with the MCP server, but not the more general JSON-RPC app
|
||||
/// server. Prefer [`OutgoingMessageSender::send_server_notification`] where
|
||||
/// possible.
|
||||
pub(crate) async fn send_event_as_notification(
|
||||
&self,
|
||||
event: &Event,
|
||||
meta: Option<OutgoingNotificationMeta>,
|
||||
) {
|
||||
#[expect(clippy::expect_used)]
|
||||
let event_json = serde_json::to_value(event).expect("Event must serialize");
|
||||
|
||||
let params = if let Ok(params) = serde_json::to_value(OutgoingNotificationParams {
|
||||
meta,
|
||||
event: event_json.clone(),
|
||||
}) {
|
||||
params
|
||||
} else {
|
||||
warn!("Failed to serialize event as OutgoingNotificationParams");
|
||||
event_json
|
||||
};
|
||||
|
||||
self.send_notification(OutgoingNotification {
|
||||
method: "codex/event".to_string(),
|
||||
params: Some(params.clone()),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn send_notification(&self, notification: OutgoingNotification) {
|
||||
let outgoing_message = OutgoingMessage::Notification(notification);
|
||||
let _ = self.sender.send(outgoing_message);
|
||||
}
|
||||
|
||||
pub(crate) async fn send_error(&self, id: RequestId, error: JSONRPCErrorError) {
|
||||
let outgoing_message = OutgoingMessage::Error(OutgoingError { id, error });
|
||||
let _ = self.sender.send(outgoing_message);
|
||||
}
|
||||
}
|
||||
|
||||
/// Outgoing message from the server to the client.
|
||||
pub(crate) enum OutgoingMessage {
|
||||
Request(OutgoingRequest),
|
||||
Notification(OutgoingNotification),
|
||||
Response(OutgoingResponse),
|
||||
Error(OutgoingError),
|
||||
}
|
||||
|
||||
impl From<OutgoingMessage> for JSONRPCMessage {
|
||||
fn from(val: OutgoingMessage) -> Self {
|
||||
use OutgoingMessage::*;
|
||||
match val {
|
||||
Request(OutgoingRequest { id, method, params }) => {
|
||||
JSONRPCMessage::Request(JSONRPCRequest {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id,
|
||||
method,
|
||||
params,
|
||||
})
|
||||
}
|
||||
Notification(OutgoingNotification { method, params }) => {
|
||||
JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
method,
|
||||
params,
|
||||
})
|
||||
}
|
||||
Response(OutgoingResponse { id, result }) => {
|
||||
JSONRPCMessage::Response(JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id,
|
||||
result,
|
||||
})
|
||||
}
|
||||
Error(OutgoingError { id, error }) => JSONRPCMessage::Error(JSONRPCError {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id,
|
||||
error,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingRequest {
|
||||
pub id: RequestId,
|
||||
pub method: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingNotification {
|
||||
pub method: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingNotificationParams {
|
||||
#[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
|
||||
pub meta: Option<OutgoingNotificationMeta>,
|
||||
|
||||
#[serde(flatten)]
|
||||
pub event: serde_json::Value,
|
||||
}
|
||||
|
||||
// Additional mcp-specific data to be added to a [`codex_core::protocol::Event`] as notification.params._meta
|
||||
// MCP Spec: https://modelcontextprotocol.io/specification/2025-06-18/basic#meta
|
||||
// Typescript Schema: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/0695a497eb50a804fc0e88c18a93a21a675d6b3e/schema/2025-06-18/schema.ts
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct OutgoingNotificationMeta {
|
||||
pub request_id: Option<RequestId>,
|
||||
}
|
||||
|
||||
impl OutgoingNotificationMeta {
|
||||
pub(crate) fn new(request_id: Option<RequestId>) -> Self {
|
||||
Self { request_id }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingResponse {
|
||||
pub id: RequestId,
|
||||
pub result: Result,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingError {
|
||||
pub error: JSONRPCErrorError,
|
||||
pub id: RequestId,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use anyhow::Result;
|
||||
use codex_core::protocol::EventMsg;
|
||||
use codex_core::protocol::SessionConfiguredEvent;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::config_types::ReasoningEffort;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_send_event_as_notification() -> Result<()> {
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
|
||||
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
|
||||
|
||||
let conversation_id = ConversationId::new();
|
||||
let rollout_file = NamedTempFile::new()?;
|
||||
let event = Event {
|
||||
id: "1".to_string(),
|
||||
msg: EventMsg::SessionConfigured(SessionConfiguredEvent {
|
||||
session_id: conversation_id,
|
||||
model: "gpt-4o".to_string(),
|
||||
reasoning_effort: Some(ReasoningEffort::default()),
|
||||
history_log_id: 1,
|
||||
history_entry_count: 1000,
|
||||
initial_messages: None,
|
||||
rollout_path: rollout_file.path().to_path_buf(),
|
||||
}),
|
||||
};
|
||||
|
||||
outgoing_message_sender
|
||||
.send_event_as_notification(&event, None)
|
||||
.await;
|
||||
|
||||
let result = outgoing_rx.recv().await.unwrap();
|
||||
let OutgoingMessage::Notification(OutgoingNotification { method, params }) = result else {
|
||||
panic!("expected Notification for first message");
|
||||
};
|
||||
assert_eq!(method, "codex/event");
|
||||
|
||||
let Ok(expected_params) = serde_json::to_value(&event) else {
|
||||
panic!("Event must serialize");
|
||||
};
|
||||
assert_eq!(params, Some(expected_params));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_send_event_as_notification_with_meta() -> Result<()> {
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
|
||||
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
|
||||
|
||||
let conversation_id = ConversationId::new();
|
||||
let rollout_file = NamedTempFile::new()?;
|
||||
let session_configured_event = SessionConfiguredEvent {
|
||||
session_id: conversation_id,
|
||||
model: "gpt-4o".to_string(),
|
||||
reasoning_effort: Some(ReasoningEffort::default()),
|
||||
history_log_id: 1,
|
||||
history_entry_count: 1000,
|
||||
initial_messages: None,
|
||||
rollout_path: rollout_file.path().to_path_buf(),
|
||||
};
|
||||
let event = Event {
|
||||
id: "1".to_string(),
|
||||
msg: EventMsg::SessionConfigured(session_configured_event.clone()),
|
||||
};
|
||||
let meta = OutgoingNotificationMeta {
|
||||
request_id: Some(RequestId::String("123".to_string())),
|
||||
};
|
||||
|
||||
outgoing_message_sender
|
||||
.send_event_as_notification(&event, Some(meta))
|
||||
.await;
|
||||
|
||||
let result = outgoing_rx.recv().await.unwrap();
|
||||
let OutgoingMessage::Notification(OutgoingNotification { method, params }) = result else {
|
||||
panic!("expected Notification for first message");
|
||||
};
|
||||
assert_eq!(method, "codex/event");
|
||||
let expected_params = json!({
|
||||
"_meta": {
|
||||
"requestId": "123",
|
||||
},
|
||||
"id": "1",
|
||||
"msg": {
|
||||
"session_id": session_configured_event.session_id,
|
||||
"model": session_configured_event.model,
|
||||
"reasoning_effort": session_configured_event.reasoning_effort,
|
||||
"history_log_id": session_configured_event.history_log_id,
|
||||
"history_entry_count": session_configured_event.history_entry_count,
|
||||
"type": "session_configured",
|
||||
"rollout_path": rollout_file.path().to_path_buf(),
|
||||
}
|
||||
});
|
||||
assert_eq!(params.unwrap(), expected_params);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
150
llmx-rs/mcp-server/src/patch_approval.rs
Normal file
150
llmx-rs/mcp-server/src/patch_approval.rs
Normal file
@@ -0,0 +1,150 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_core::CodexConversation;
|
||||
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,
|
||||
pub codex_call_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(
|
||||
call_id: String,
|
||||
reason: Option<String>,
|
||||
grant_root: Option<PathBuf>,
|
||||
changes: HashMap<PathBuf, FileChange>,
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
codex: Arc<CodexConversation>,
|
||||
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_call_id: call_id,
|
||||
codex_reason: reason,
|
||||
codex_grant_root: grant_root,
|
||||
codex_changes: changes,
|
||||
};
|
||||
let params_json = match serde_json::to_value(¶ms) {
|
||||
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<CodexConversation>,
|
||||
) {
|
||||
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}");
|
||||
}
|
||||
}
|
||||
2
llmx-rs/mcp-server/src/tool_handlers/mod.rs
Normal file
2
llmx-rs/mcp-server/src/tool_handlers/mod.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
pub(crate) mod create_conversation;
|
||||
pub(crate) mod send_message;
|
||||
3
llmx-rs/mcp-server/tests/all.rs
Normal file
3
llmx-rs/mcp-server/tests/all.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
// Single integration test binary that aggregates all test modules.
|
||||
// The submodules live in `tests/suite/`.
|
||||
mod suite;
|
||||
25
llmx-rs/mcp-server/tests/common/Cargo.toml
Normal file
25
llmx-rs/mcp-server/tests/common/Cargo.toml
Normal file
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
edition = "2024"
|
||||
name = "mcp_test_support"
|
||||
version = { workspace = true }
|
||||
|
||||
[lib]
|
||||
path = "lib.rs"
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
assert_cmd = { workspace = true }
|
||||
codex-core = { workspace = true }
|
||||
codex-mcp-server = { workspace = true }
|
||||
mcp-types = { workspace = true }
|
||||
os_info = { workspace = true }
|
||||
pretty_assertions = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tokio = { workspace = true, features = [
|
||||
"io-std",
|
||||
"macros",
|
||||
"process",
|
||||
"rt-multi-thread",
|
||||
] }
|
||||
wiremock = { workspace = true }
|
||||
17
llmx-rs/mcp-server/tests/common/lib.rs
Normal file
17
llmx-rs/mcp-server/tests/common/lib.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
mod mcp_process;
|
||||
mod mock_model_server;
|
||||
mod responses;
|
||||
|
||||
pub use mcp_process::McpProcess;
|
||||
use mcp_types::JSONRPCResponse;
|
||||
pub use mock_model_server::create_mock_chat_completions_server;
|
||||
pub use responses::create_apply_patch_sse_response;
|
||||
pub use responses::create_final_assistant_message_sse_response;
|
||||
pub use responses::create_shell_sse_response;
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
pub fn to_response<T: DeserializeOwned>(response: JSONRPCResponse) -> anyhow::Result<T> {
|
||||
let value = serde_json::to_value(response.result)?;
|
||||
let codex_response = serde_json::from_value(value)?;
|
||||
Ok(codex_response)
|
||||
}
|
||||
343
llmx-rs/mcp-server/tests/common/mcp_process.rs
Normal file
343
llmx-rs/mcp-server/tests/common/mcp_process.rs
Normal file
@@ -0,0 +1,343 @@
|
||||
use std::path::Path;
|
||||
use std::process::Stdio;
|
||||
use std::sync::atomic::AtomicI64;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::process::Child;
|
||||
use tokio::process::ChildStdin;
|
||||
use tokio::process::ChildStdout;
|
||||
|
||||
use anyhow::Context;
|
||||
use assert_cmd::prelude::*;
|
||||
use codex_mcp_server::CodexToolCallParam;
|
||||
|
||||
use mcp_types::CallToolRequestParams;
|
||||
use mcp_types::ClientCapabilities;
|
||||
use mcp_types::Implementation;
|
||||
use mcp_types::InitializeRequestParams;
|
||||
use mcp_types::JSONRPC_VERSION;
|
||||
use mcp_types::JSONRPCMessage;
|
||||
use mcp_types::JSONRPCNotification;
|
||||
use mcp_types::JSONRPCRequest;
|
||||
use mcp_types::JSONRPCResponse;
|
||||
use mcp_types::ModelContextProtocolNotification;
|
||||
use mcp_types::ModelContextProtocolRequest;
|
||||
use mcp_types::RequestId;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::process::Command as StdCommand;
|
||||
use tokio::process::Command;
|
||||
|
||||
pub struct McpProcess {
|
||||
next_request_id: AtomicI64,
|
||||
/// Retain this child process until the client is dropped. The Tokio runtime
|
||||
/// will make a "best effort" to reap the process after it exits, but it is
|
||||
/// not a guarantee. See the `kill_on_drop` documentation for details.
|
||||
#[allow(dead_code)]
|
||||
process: Child,
|
||||
stdin: ChildStdin,
|
||||
stdout: BufReader<ChildStdout>,
|
||||
}
|
||||
|
||||
impl McpProcess {
|
||||
pub async fn new(codex_home: &Path) -> anyhow::Result<Self> {
|
||||
Self::new_with_env(codex_home, &[]).await
|
||||
}
|
||||
|
||||
/// Creates a new MCP process, allowing tests to override or remove
|
||||
/// specific environment variables for the child process only.
|
||||
///
|
||||
/// Pass a tuple of (key, Some(value)) to set/override, or (key, None) to
|
||||
/// remove a variable from the child's environment.
|
||||
pub async fn new_with_env(
|
||||
codex_home: &Path,
|
||||
env_overrides: &[(&str, Option<&str>)],
|
||||
) -> anyhow::Result<Self> {
|
||||
// Use assert_cmd to locate the binary path and then switch to tokio::process::Command
|
||||
let std_cmd = StdCommand::cargo_bin("codex-mcp-server")
|
||||
.context("should find binary for codex-mcp-server")?;
|
||||
|
||||
let program = std_cmd.get_program().to_owned();
|
||||
|
||||
let mut cmd = Command::new(program);
|
||||
|
||||
cmd.stdin(Stdio::piped());
|
||||
cmd.stdout(Stdio::piped());
|
||||
cmd.stderr(Stdio::piped());
|
||||
cmd.env("CODEX_HOME", codex_home);
|
||||
cmd.env("RUST_LOG", "debug");
|
||||
|
||||
for (k, v) in env_overrides {
|
||||
match v {
|
||||
Some(val) => {
|
||||
cmd.env(k, val);
|
||||
}
|
||||
None => {
|
||||
cmd.env_remove(k);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut process = cmd
|
||||
.kill_on_drop(true)
|
||||
.spawn()
|
||||
.context("codex-mcp-server proc should start")?;
|
||||
let stdin = process
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| anyhow::format_err!("mcp should have stdin fd"))?;
|
||||
let stdout = process
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| anyhow::format_err!("mcp should have stdout fd"))?;
|
||||
let stdout = BufReader::new(stdout);
|
||||
|
||||
// Forward child's stderr to our stderr so failures are visible even
|
||||
// when stdout/stderr are captured by the test harness.
|
||||
if let Some(stderr) = process.stderr.take() {
|
||||
let mut stderr_reader = BufReader::new(stderr).lines();
|
||||
tokio::spawn(async move {
|
||||
while let Ok(Some(line)) = stderr_reader.next_line().await {
|
||||
eprintln!("[mcp stderr] {line}");
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
next_request_id: AtomicI64::new(0),
|
||||
process,
|
||||
stdin,
|
||||
stdout,
|
||||
})
|
||||
}
|
||||
|
||||
/// Performs the initialization handshake with the MCP server.
|
||||
pub async fn initialize(&mut self) -> anyhow::Result<()> {
|
||||
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
let params = InitializeRequestParams {
|
||||
capabilities: ClientCapabilities {
|
||||
elicitation: Some(json!({})),
|
||||
experimental: None,
|
||||
roots: None,
|
||||
sampling: None,
|
||||
},
|
||||
client_info: Implementation {
|
||||
name: "elicitation test".into(),
|
||||
title: Some("Elicitation Test".into()),
|
||||
version: "0.0.0".into(),
|
||||
user_agent: None,
|
||||
},
|
||||
protocol_version: mcp_types::MCP_SCHEMA_VERSION.into(),
|
||||
};
|
||||
let params_value = serde_json::to_value(params)?;
|
||||
|
||||
self.send_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(request_id),
|
||||
method: mcp_types::InitializeRequest::METHOD.into(),
|
||||
params: Some(params_value),
|
||||
}))
|
||||
.await?;
|
||||
|
||||
let initialized = self.read_jsonrpc_message().await?;
|
||||
let os_info = os_info::get();
|
||||
let user_agent = format!(
|
||||
"codex_cli_rs/0.0.0 ({} {}; {}) {} (elicitation test; 0.0.0)",
|
||||
os_info.os_type(),
|
||||
os_info.version(),
|
||||
os_info.architecture().unwrap_or("unknown"),
|
||||
codex_core::terminal::user_agent()
|
||||
);
|
||||
assert_eq!(
|
||||
JSONRPCMessage::Response(JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(request_id),
|
||||
result: json!({
|
||||
"capabilities": {
|
||||
"tools": {
|
||||
"listChanged": true
|
||||
},
|
||||
},
|
||||
"serverInfo": {
|
||||
"name": "codex-mcp-server",
|
||||
"title": "Codex",
|
||||
"version": "0.0.0",
|
||||
"user_agent": user_agent
|
||||
},
|
||||
"protocolVersion": mcp_types::MCP_SCHEMA_VERSION
|
||||
})
|
||||
}),
|
||||
initialized
|
||||
);
|
||||
|
||||
// Send notifications/initialized to ack the response.
|
||||
self.send_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
method: mcp_types::InitializedNotification::METHOD.into(),
|
||||
params: None,
|
||||
}))
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the id used to make the request so it can be used when
|
||||
/// correlating notifications.
|
||||
pub async fn send_codex_tool_call(
|
||||
&mut self,
|
||||
params: CodexToolCallParam,
|
||||
) -> anyhow::Result<i64> {
|
||||
let codex_tool_call_params = CallToolRequestParams {
|
||||
name: "codex".to_string(),
|
||||
arguments: Some(serde_json::to_value(params)?),
|
||||
};
|
||||
self.send_request(
|
||||
mcp_types::CallToolRequest::METHOD,
|
||||
Some(serde_json::to_value(codex_tool_call_params)?),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_request(
|
||||
&mut self,
|
||||
method: &str,
|
||||
params: Option<serde_json::Value>,
|
||||
) -> anyhow::Result<i64> {
|
||||
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
let message = JSONRPCMessage::Request(JSONRPCRequest {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(request_id),
|
||||
method: method.to_string(),
|
||||
params,
|
||||
});
|
||||
self.send_jsonrpc_message(message).await?;
|
||||
Ok(request_id)
|
||||
}
|
||||
|
||||
pub async fn send_response(
|
||||
&mut self,
|
||||
id: RequestId,
|
||||
result: serde_json::Value,
|
||||
) -> anyhow::Result<()> {
|
||||
self.send_jsonrpc_message(JSONRPCMessage::Response(JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id,
|
||||
result,
|
||||
}))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_jsonrpc_message(&mut self, message: JSONRPCMessage) -> anyhow::Result<()> {
|
||||
eprintln!("writing message to stdin: {message:?}");
|
||||
let payload = serde_json::to_string(&message)?;
|
||||
self.stdin.write_all(payload.as_bytes()).await?;
|
||||
self.stdin.write_all(b"\n").await?;
|
||||
self.stdin.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_jsonrpc_message(&mut self) -> anyhow::Result<JSONRPCMessage> {
|
||||
let mut line = String::new();
|
||||
self.stdout.read_line(&mut line).await?;
|
||||
let message = serde_json::from_str::<JSONRPCMessage>(&line)?;
|
||||
eprintln!("read message from stdout: {message:?}");
|
||||
Ok(message)
|
||||
}
|
||||
|
||||
pub async fn read_stream_until_request_message(&mut self) -> anyhow::Result<JSONRPCRequest> {
|
||||
eprintln!("in read_stream_until_request_message()");
|
||||
|
||||
loop {
|
||||
let message = self.read_jsonrpc_message().await?;
|
||||
|
||||
match message {
|
||||
JSONRPCMessage::Notification(_) => {
|
||||
eprintln!("notification: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Request(jsonrpc_request) => {
|
||||
return Ok(jsonrpc_request);
|
||||
}
|
||||
JSONRPCMessage::Error(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Error: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Response(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Response: {message:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn read_stream_until_response_message(
|
||||
&mut self,
|
||||
request_id: RequestId,
|
||||
) -> anyhow::Result<JSONRPCResponse> {
|
||||
eprintln!("in read_stream_until_response_message({request_id:?})");
|
||||
|
||||
loop {
|
||||
let message = self.read_jsonrpc_message().await?;
|
||||
match message {
|
||||
JSONRPCMessage::Notification(_) => {
|
||||
eprintln!("notification: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Request(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Request: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Error(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Error: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Response(jsonrpc_response) => {
|
||||
if jsonrpc_response.id == request_id {
|
||||
return Ok(jsonrpc_response);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads notifications until a legacy TaskComplete event is observed:
|
||||
/// Method "codex/event" with params.msg.type == "task_complete".
|
||||
pub async fn read_stream_until_legacy_task_complete_notification(
|
||||
&mut self,
|
||||
) -> anyhow::Result<JSONRPCNotification> {
|
||||
eprintln!("in read_stream_until_legacy_task_complete_notification()");
|
||||
|
||||
loop {
|
||||
let message = self.read_jsonrpc_message().await?;
|
||||
match message {
|
||||
JSONRPCMessage::Notification(notification) => {
|
||||
let is_match = if notification.method == "codex/event" {
|
||||
if let Some(params) = ¬ification.params {
|
||||
params
|
||||
.get("msg")
|
||||
.and_then(|m| m.get("type"))
|
||||
.and_then(|t| t.as_str())
|
||||
== Some("task_complete")
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if is_match {
|
||||
return Ok(notification);
|
||||
} else {
|
||||
eprintln!("ignoring notification: {notification:?}");
|
||||
}
|
||||
}
|
||||
JSONRPCMessage::Request(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Request: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Error(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Error: {message:?}");
|
||||
}
|
||||
JSONRPCMessage::Response(_) => {
|
||||
anyhow::bail!("unexpected JSONRPCMessage::Response: {message:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
47
llmx-rs/mcp-server/tests/common/mock_model_server.rs
Normal file
47
llmx-rs/mcp-server/tests/common/mock_model_server.rs
Normal file
@@ -0,0 +1,47 @@
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::Respond;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
/// Create a mock server that will provide the responses, in order, for
|
||||
/// requests to the `/v1/chat/completions` endpoint.
|
||||
pub async fn create_mock_chat_completions_server(responses: Vec<String>) -> MockServer {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let num_calls = responses.len();
|
||||
let seq_responder = SeqResponder {
|
||||
num_calls: AtomicUsize::new(0),
|
||||
responses,
|
||||
};
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(seq_responder)
|
||||
.expect(num_calls as u64)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
server
|
||||
}
|
||||
|
||||
struct SeqResponder {
|
||||
num_calls: AtomicUsize,
|
||||
responses: Vec<String>,
|
||||
}
|
||||
|
||||
impl Respond for SeqResponder {
|
||||
fn respond(&self, _: &wiremock::Request) -> ResponseTemplate {
|
||||
let call_num = self.num_calls.fetch_add(1, Ordering::SeqCst);
|
||||
match self.responses.get(call_num) {
|
||||
Some(response) => ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_raw(response.clone(), "text/event-stream"),
|
||||
None => panic!("no response for {call_num}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
95
llmx-rs/mcp-server/tests/common/responses.rs
Normal file
95
llmx-rs/mcp-server/tests/common/responses.rs
Normal file
@@ -0,0 +1,95 @@
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
|
||||
pub fn create_shell_sse_response(
|
||||
command: Vec<String>,
|
||||
workdir: Option<&Path>,
|
||||
timeout_ms: Option<u64>,
|
||||
call_id: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
// The `arguments`` for the `shell` tool is a serialized JSON object.
|
||||
let tool_call_arguments = serde_json::to_string(&json!({
|
||||
"command": command,
|
||||
"workdir": workdir.map(|w| w.to_string_lossy()),
|
||||
"timeout": timeout_ms
|
||||
}))?;
|
||||
let tool_call = json!({
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call_id,
|
||||
"function": {
|
||||
"name": "shell",
|
||||
"arguments": tool_call_arguments
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let sse = format!(
|
||||
"data: {}\n\ndata: DONE\n\n",
|
||||
serde_json::to_string(&tool_call)?
|
||||
);
|
||||
Ok(sse)
|
||||
}
|
||||
|
||||
pub fn create_final_assistant_message_sse_response(message: &str) -> anyhow::Result<String> {
|
||||
let assistant_message = json!({
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"content": message
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let sse = format!(
|
||||
"data: {}\n\ndata: DONE\n\n",
|
||||
serde_json::to_string(&assistant_message)?
|
||||
);
|
||||
Ok(sse)
|
||||
}
|
||||
|
||||
pub fn create_apply_patch_sse_response(
|
||||
patch_content: &str,
|
||||
call_id: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
// Use shell command to call apply_patch with heredoc format
|
||||
let shell_command = format!("apply_patch <<'EOF'\n{patch_content}\nEOF");
|
||||
let tool_call_arguments = serde_json::to_string(&json!({
|
||||
"command": ["bash", "-lc", shell_command]
|
||||
}))?;
|
||||
|
||||
let tool_call = json!({
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call_id,
|
||||
"function": {
|
||||
"name": "shell",
|
||||
"arguments": tool_call_arguments
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let sse = format!(
|
||||
"data: {}\n\ndata: DONE\n\n",
|
||||
serde_json::to_string(&tool_call)?
|
||||
);
|
||||
Ok(sse)
|
||||
}
|
||||
481
llmx-rs/mcp-server/tests/suite/codex_tool.rs
Normal file
481
llmx-rs/mcp-server/tests/suite/codex_tool.rs
Normal file
@@ -0,0 +1,481 @@
|
||||
use std::collections::HashMap;
|
||||
use std::env;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use codex_core::parse_command;
|
||||
use codex_core::protocol::FileChange;
|
||||
use codex_core::protocol::ReviewDecision;
|
||||
use codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR;
|
||||
use codex_mcp_server::CodexToolCallParam;
|
||||
use codex_mcp_server::ExecApprovalElicitRequestParams;
|
||||
use codex_mcp_server::ExecApprovalResponse;
|
||||
use codex_mcp_server::PatchApprovalElicitRequestParams;
|
||||
use codex_mcp_server::PatchApprovalResponse;
|
||||
use mcp_types::ElicitRequest;
|
||||
use mcp_types::ElicitRequestParamsRequestedSchema;
|
||||
use mcp_types::JSONRPC_VERSION;
|
||||
use mcp_types::JSONRPCRequest;
|
||||
use mcp_types::JSONRPCResponse;
|
||||
use mcp_types::ModelContextProtocolRequest;
|
||||
use mcp_types::RequestId;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tokio::time::timeout;
|
||||
use wiremock::MockServer;
|
||||
|
||||
use core_test_support::skip_if_no_network;
|
||||
use mcp_test_support::McpProcess;
|
||||
use mcp_test_support::create_apply_patch_sse_response;
|
||||
use mcp_test_support::create_final_assistant_message_sse_response;
|
||||
use mcp_test_support::create_mock_chat_completions_server;
|
||||
use mcp_test_support::create_shell_sse_response;
|
||||
|
||||
// Allow ample time on slower CI or under load to avoid flakes.
|
||||
const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
|
||||
|
||||
/// Test that a shell command that is not on the "trusted" list triggers an
|
||||
/// elicitation request to the MCP and that sending the approval runs the
|
||||
/// command, as expected.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn test_shell_command_approval_triggers_elicitation() {
|
||||
if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
|
||||
println!(
|
||||
"Skipping test because it cannot execute when network is disabled in a Codex sandbox."
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Apparently `#[tokio::test]` must return `()`, so we create a helper
|
||||
// function that returns `Result` so we can use `?` in favor of `unwrap`.
|
||||
if let Err(err) = shell_command_approval_triggers_elicitation().await {
|
||||
panic!("failure: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
async fn shell_command_approval_triggers_elicitation() -> anyhow::Result<()> {
|
||||
// Use a simple, untrusted command that creates a file so we can
|
||||
// observe a side-effect.
|
||||
//
|
||||
// Cross‑platform approach: run a tiny Python snippet to touch the file
|
||||
// using `python3 -c ...` on all platforms.
|
||||
let workdir_for_shell_function_call = TempDir::new()?;
|
||||
let created_filename = "created_by_shell_tool.txt";
|
||||
let created_file = workdir_for_shell_function_call
|
||||
.path()
|
||||
.join(created_filename);
|
||||
|
||||
let shell_command = vec![
|
||||
"python3".to_string(),
|
||||
"-c".to_string(),
|
||||
format!("import pathlib; pathlib.Path('{created_filename}').touch()"),
|
||||
];
|
||||
|
||||
let McpHandle {
|
||||
process: mut mcp_process,
|
||||
server: _server,
|
||||
dir: _dir,
|
||||
} = create_mcp_process(vec![
|
||||
create_shell_sse_response(
|
||||
shell_command.clone(),
|
||||
Some(workdir_for_shell_function_call.path()),
|
||||
Some(5_000),
|
||||
"call1234",
|
||||
)?,
|
||||
create_final_assistant_message_sse_response("File created!")?,
|
||||
])
|
||||
.await?;
|
||||
|
||||
// Send a "codex" tool request, which should hit the completions endpoint.
|
||||
// In turn, it should reply with a tool call, which the MCP should forward
|
||||
// as an elicitation.
|
||||
let codex_request_id = mcp_process
|
||||
.send_codex_tool_call(CodexToolCallParam {
|
||||
prompt: "run `git init`".to_string(),
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let elicitation_request = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_request_message(),
|
||||
)
|
||||
.await??;
|
||||
|
||||
let elicitation_request_id = elicitation_request.id.clone();
|
||||
let params = serde_json::from_value::<ExecApprovalElicitRequestParams>(
|
||||
elicitation_request
|
||||
.params
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow::anyhow!("elicitation_request.params must be set"))?,
|
||||
)?;
|
||||
let expected_elicitation_request = create_expected_elicitation_request(
|
||||
elicitation_request_id.clone(),
|
||||
shell_command.clone(),
|
||||
workdir_for_shell_function_call.path(),
|
||||
codex_request_id.to_string(),
|
||||
params.codex_event_id.clone(),
|
||||
)?;
|
||||
assert_eq!(expected_elicitation_request, elicitation_request);
|
||||
|
||||
// Accept the `git init` request by responding to the elicitation.
|
||||
mcp_process
|
||||
.send_response(
|
||||
elicitation_request_id,
|
||||
serde_json::to_value(ExecApprovalResponse {
|
||||
decision: ReviewDecision::Approved,
|
||||
})?,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Verify task_complete notification arrives before the tool call completes.
|
||||
#[expect(clippy::expect_used)]
|
||||
let _task_complete = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_legacy_task_complete_notification(),
|
||||
)
|
||||
.await
|
||||
.expect("task_complete_notification timeout")
|
||||
.expect("task_complete_notification resp");
|
||||
|
||||
// Verify the original `codex` tool call completes and that the file was created.
|
||||
let codex_response = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_response_message(RequestId::Integer(codex_request_id)),
|
||||
)
|
||||
.await??;
|
||||
assert_eq!(
|
||||
JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(codex_request_id),
|
||||
result: json!({
|
||||
"content": [
|
||||
{
|
||||
"text": "File created!",
|
||||
"type": "text"
|
||||
}
|
||||
]
|
||||
}),
|
||||
},
|
||||
codex_response
|
||||
);
|
||||
|
||||
assert!(created_file.is_file(), "created file should exist");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_expected_elicitation_request(
|
||||
elicitation_request_id: RequestId,
|
||||
command: Vec<String>,
|
||||
workdir: &Path,
|
||||
codex_mcp_tool_call_id: String,
|
||||
codex_event_id: String,
|
||||
) -> anyhow::Result<JSONRPCRequest> {
|
||||
let expected_message = format!(
|
||||
"Allow Codex to run `{}` in `{}`?",
|
||||
shlex::try_join(command.iter().map(std::convert::AsRef::as_ref))?,
|
||||
workdir.to_string_lossy()
|
||||
);
|
||||
let codex_parsed_cmd = parse_command::parse_command(&command);
|
||||
Ok(JSONRPCRequest {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: elicitation_request_id,
|
||||
method: ElicitRequest::METHOD.to_string(),
|
||||
params: Some(serde_json::to_value(&ExecApprovalElicitRequestParams {
|
||||
message: expected_message,
|
||||
requested_schema: ElicitRequestParamsRequestedSchema {
|
||||
r#type: "object".to_string(),
|
||||
properties: json!({}),
|
||||
required: None,
|
||||
},
|
||||
codex_elicitation: "exec-approval".to_string(),
|
||||
codex_mcp_tool_call_id,
|
||||
codex_event_id,
|
||||
codex_command: command,
|
||||
codex_cwd: workdir.to_path_buf(),
|
||||
codex_call_id: "call1234".to_string(),
|
||||
codex_parsed_cmd,
|
||||
codex_risk: None,
|
||||
})?),
|
||||
})
|
||||
}
|
||||
|
||||
/// Test that patch approval triggers an elicitation request to the MCP and that
|
||||
/// sending the approval applies the patch, as expected.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_patch_approval_triggers_elicitation() {
|
||||
if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
|
||||
println!(
|
||||
"Skipping test because it cannot execute when network is disabled in a Codex sandbox."
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Err(err) = patch_approval_triggers_elicitation().await {
|
||||
panic!("failure: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
async fn patch_approval_triggers_elicitation() -> anyhow::Result<()> {
|
||||
let cwd = TempDir::new()?;
|
||||
let test_file = cwd.path().join("destination_file.txt");
|
||||
std::fs::write(&test_file, "original content\n")?;
|
||||
|
||||
let patch_content = format!(
|
||||
"*** Begin Patch\n*** Update File: {}\n-original content\n+modified content\n*** End Patch",
|
||||
test_file.as_path().to_string_lossy()
|
||||
);
|
||||
|
||||
let McpHandle {
|
||||
process: mut mcp_process,
|
||||
server: _server,
|
||||
dir: _dir,
|
||||
} = create_mcp_process(vec![
|
||||
create_apply_patch_sse_response(&patch_content, "call1234")?,
|
||||
create_final_assistant_message_sse_response("Patch has been applied successfully!")?,
|
||||
])
|
||||
.await?;
|
||||
|
||||
// Send a "codex" tool request that will trigger the apply_patch command
|
||||
let codex_request_id = mcp_process
|
||||
.send_codex_tool_call(CodexToolCallParam {
|
||||
cwd: Some(cwd.path().to_string_lossy().to_string()),
|
||||
prompt: "please modify the test file".to_string(),
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let elicitation_request = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_request_message(),
|
||||
)
|
||||
.await??;
|
||||
|
||||
let elicitation_request_id = RequestId::Integer(0);
|
||||
|
||||
let mut expected_changes = HashMap::new();
|
||||
expected_changes.insert(
|
||||
test_file.as_path().to_path_buf(),
|
||||
FileChange::Update {
|
||||
unified_diff: "@@ -1 +1 @@\n-original content\n+modified content\n".to_string(),
|
||||
move_path: None,
|
||||
},
|
||||
);
|
||||
|
||||
let expected_elicitation_request = create_expected_patch_approval_elicitation_request(
|
||||
elicitation_request_id.clone(),
|
||||
expected_changes,
|
||||
None, // No grant_root expected
|
||||
None, // No reason expected
|
||||
codex_request_id.to_string(),
|
||||
"1".to_string(),
|
||||
)?;
|
||||
assert_eq!(expected_elicitation_request, elicitation_request);
|
||||
|
||||
// Accept the patch approval request by responding to the elicitation
|
||||
mcp_process
|
||||
.send_response(
|
||||
elicitation_request_id,
|
||||
serde_json::to_value(PatchApprovalResponse {
|
||||
decision: ReviewDecision::Approved,
|
||||
})?,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Verify the original `codex` tool call completes
|
||||
let codex_response = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_response_message(RequestId::Integer(codex_request_id)),
|
||||
)
|
||||
.await??;
|
||||
assert_eq!(
|
||||
JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(codex_request_id),
|
||||
result: json!({
|
||||
"content": [
|
||||
{
|
||||
"text": "Patch has been applied successfully!",
|
||||
"type": "text"
|
||||
}
|
||||
]
|
||||
}),
|
||||
},
|
||||
codex_response
|
||||
);
|
||||
|
||||
let file_contents = std::fs::read_to_string(test_file.as_path())?;
|
||||
assert_eq!(file_contents, "modified content\n");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_codex_tool_passes_base_instructions() {
|
||||
skip_if_no_network!();
|
||||
|
||||
// Apparently `#[tokio::test]` must return `()`, so we create a helper
|
||||
// function that returns `Result` so we can use `?` in favor of `unwrap`.
|
||||
if let Err(err) = codex_tool_passes_base_instructions().await {
|
||||
panic!("failure: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
async fn codex_tool_passes_base_instructions() -> anyhow::Result<()> {
|
||||
#![expect(clippy::unwrap_used)]
|
||||
|
||||
let server =
|
||||
create_mock_chat_completions_server(vec![create_final_assistant_message_sse_response(
|
||||
"Enjoy!",
|
||||
)?])
|
||||
.await;
|
||||
|
||||
// Run `codex mcp` with a specific config.toml.
|
||||
let codex_home = TempDir::new()?;
|
||||
create_config_toml(codex_home.path(), &server.uri())?;
|
||||
let mut mcp_process = McpProcess::new(codex_home.path()).await?;
|
||||
timeout(DEFAULT_READ_TIMEOUT, mcp_process.initialize()).await??;
|
||||
|
||||
// Send a "codex" tool request, which should hit the completions endpoint.
|
||||
let codex_request_id = mcp_process
|
||||
.send_codex_tool_call(CodexToolCallParam {
|
||||
prompt: "How are you?".to_string(),
|
||||
base_instructions: Some("You are a helpful assistant.".to_string()),
|
||||
developer_instructions: Some("Foreshadow upcoming tool calls.".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
|
||||
let codex_response = timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp_process.read_stream_until_response_message(RequestId::Integer(codex_request_id)),
|
||||
)
|
||||
.await??;
|
||||
assert_eq!(
|
||||
JSONRPCResponse {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: RequestId::Integer(codex_request_id),
|
||||
result: json!({
|
||||
"content": [
|
||||
{
|
||||
"text": "Enjoy!",
|
||||
"type": "text"
|
||||
}
|
||||
]
|
||||
}),
|
||||
},
|
||||
codex_response
|
||||
);
|
||||
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let request = requests[0].body_json::<serde_json::Value>()?;
|
||||
let instructions = request["messages"][0]["content"].as_str().unwrap();
|
||||
assert!(instructions.starts_with("You are a helpful assistant."));
|
||||
|
||||
let developer_msg = request["messages"]
|
||||
.as_array()
|
||||
.and_then(|messages| {
|
||||
messages
|
||||
.iter()
|
||||
.find(|msg| msg.get("role").and_then(|role| role.as_str()) == Some("developer"))
|
||||
})
|
||||
.unwrap();
|
||||
let developer_content = developer_msg
|
||||
.get("content")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap();
|
||||
assert!(
|
||||
!developer_content.contains('<'),
|
||||
"expected developer instructions without XML tags, got `{developer_content}`"
|
||||
);
|
||||
assert_eq!(developer_content, "Foreshadow upcoming tool calls.");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_expected_patch_approval_elicitation_request(
|
||||
elicitation_request_id: RequestId,
|
||||
changes: HashMap<PathBuf, FileChange>,
|
||||
grant_root: Option<PathBuf>,
|
||||
reason: Option<String>,
|
||||
codex_mcp_tool_call_id: String,
|
||||
codex_event_id: String,
|
||||
) -> anyhow::Result<JSONRPCRequest> {
|
||||
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());
|
||||
|
||||
Ok(JSONRPCRequest {
|
||||
jsonrpc: JSONRPC_VERSION.into(),
|
||||
id: elicitation_request_id,
|
||||
method: ElicitRequest::METHOD.to_string(),
|
||||
params: Some(serde_json::to_value(&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,
|
||||
codex_event_id,
|
||||
codex_reason: reason,
|
||||
codex_grant_root: grant_root,
|
||||
codex_changes: changes,
|
||||
codex_call_id: "call1234".to_string(),
|
||||
})?),
|
||||
})
|
||||
}
|
||||
|
||||
/// This handle is used to ensure that the MockServer and TempDir are not dropped while
|
||||
/// the McpProcess is still running.
|
||||
pub struct McpHandle {
|
||||
pub process: McpProcess,
|
||||
/// Retain the server for the lifetime of the McpProcess.
|
||||
#[allow(dead_code)]
|
||||
server: MockServer,
|
||||
/// Retain the temporary directory for the lifetime of the McpProcess.
|
||||
#[allow(dead_code)]
|
||||
dir: TempDir,
|
||||
}
|
||||
|
||||
async fn create_mcp_process(responses: Vec<String>) -> anyhow::Result<McpHandle> {
|
||||
let server = create_mock_chat_completions_server(responses).await;
|
||||
let codex_home = TempDir::new()?;
|
||||
create_config_toml(codex_home.path(), &server.uri())?;
|
||||
let mut mcp_process = McpProcess::new(codex_home.path()).await?;
|
||||
timeout(DEFAULT_READ_TIMEOUT, mcp_process.initialize()).await??;
|
||||
Ok(McpHandle {
|
||||
process: mcp_process,
|
||||
server,
|
||||
dir: codex_home,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a Codex config that uses the mock server as the model provider.
|
||||
/// It also uses `approval_policy = "untrusted"` so that we exercise the
|
||||
/// elicitation code path for shell commands.
|
||||
fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {
|
||||
let config_toml = codex_home.join("config.toml");
|
||||
std::fs::write(
|
||||
config_toml,
|
||||
format!(
|
||||
r#"
|
||||
model = "mock-model"
|
||||
approval_policy = "untrusted"
|
||||
sandbox_policy = "workspace-write"
|
||||
|
||||
model_provider = "mock_provider"
|
||||
|
||||
[model_providers.mock_provider]
|
||||
name = "Mock provider for test"
|
||||
base_url = "{server_uri}/v1"
|
||||
wire_api = "chat"
|
||||
request_max_retries = 0
|
||||
stream_max_retries = 0
|
||||
"#
|
||||
),
|
||||
)
|
||||
}
|
||||
1
llmx-rs/mcp-server/tests/suite/mod.rs
Normal file
1
llmx-rs/mcp-server/tests/suite/mod.rs
Normal file
@@ -0,0 +1 @@
|
||||
mod codex_tool;
|
||||
Reference in New Issue
Block a user