Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions codex-rs/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

86 changes: 84 additions & 2 deletions codex-rs/core/src/default_client.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,17 @@
use reqwest::header::HeaderValue;
use std::sync::LazyLock;
use std::sync::Mutex;

/// Set this to add a suffix to the User-Agent string.
///
/// This is primarily designed to differentiate MCP clients (or the lack thereof) from each other.
/// Because there can only be one MCP server per process, it should be safe for this to be a global static.
/// However, future users of this should use this with caution as a result.
///
/// A space is automatically added between the suffix and the rest of the User-Agent string.
/// The full user agent string is returned from the mcp initialize response.
/// Parenthesis will be added by Codex. This should only specify what goes inside of the parenthesis.
pub static USER_AGENT_SUFFIX: LazyLock<Mutex<Option<String>>> = LazyLock::new(|| Mutex::new(None));

pub const CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR: &str = "CODEX_INTERNAL_ORIGINATOR_OVERRIDE";

Expand Down Expand Up @@ -32,14 +44,62 @@ pub static ORIGINATOR: LazyLock<Originator> = LazyLock::new(|| {
pub fn get_codex_user_agent() -> String {
let build_version = env!("CARGO_PKG_VERSION");
let os_info = os_info::get();
format!(
let prefix = format!(
"{}/{build_version} ({} {}; {}) {}",
ORIGINATOR.value.as_str(),
os_info.os_type(),
os_info.version(),
os_info.architecture().unwrap_or("unknown"),
crate::terminal::user_agent()
)
);
let suffix = USER_AGENT_SUFFIX
.lock()
.ok()
.and_then(|guard| guard.clone())
.and_then(|value| {
let value = value.trim();
if value.is_empty() {
None
} else {
Some(value.to_string())
}
})
.map_or_else(String::new, |value| format!(" ({value})"));

let candidate = format!("{prefix}{suffix}");
sanitize_user_agent(candidate, &prefix)
}

/// Sanitize the user agent string.
///
/// Invalid characters are replaced with an underscore.
///
/// If the user agent fails to parse, it falls back to fallback and then to ORIGINATOR.
fn sanitize_user_agent(candidate: String, fallback: &str) -> String {
if HeaderValue::from_str(candidate.as_str()).is_ok() {
return candidate;
}

let sanitized: String = candidate
.chars()
.map(|ch| if matches!(ch, ' '..='~') { ch } else { '_' })
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we sure we want to allow symbols and spaces? Can't hurt, I guess?

.collect();
if !sanitized.is_empty() && HeaderValue::from_str(sanitized.as_str()).is_ok() {
tracing::warn!(
"Sanitized Codex user agent because provided suffix contained invalid header characters"
);
sanitized
} else if HeaderValue::from_str(fallback).is_ok() {
tracing::warn!(
"Falling back to base Codex user agent because provided suffix could not be sanitized"
);
fallback.to_string()
} else {
tracing::warn!(
"Falling back to default Codex originator because base user agent string is invalid"
);
ORIGINATOR.value.clone()
}
}

/// Create a reqwest client with default `originator` and `User-Agent` headers set.
Expand Down Expand Up @@ -114,6 +174,28 @@ mod tests {
assert_eq!(ua_header.to_str().unwrap(), expected_ua);
}

#[test]
fn test_invalid_suffix_is_sanitized() {
let prefix = "codex_cli_rs/0.0.0";
let suffix = "bad\rsuffix";

assert_eq!(
sanitize_user_agent(format!("{prefix} ({suffix})"), prefix),
"codex_cli_rs/0.0.0 (bad_suffix)"
);
}

#[test]
fn test_invalid_suffix_is_sanitized2() {
let prefix = "codex_cli_rs/0.0.0";
let suffix = "bad\0suffix";

assert_eq!(
sanitize_user_agent(format!("{prefix} ({suffix})"), prefix),
"codex_cli_rs/0.0.0 (bad_suffix)"
);
}

#[test]
#[cfg(target_os = "macos")]
fn test_macos() {
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/core/src/mcp_connection_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use anyhow::Result;
use anyhow::anyhow;
use codex_mcp_client::McpClient;
use mcp_types::ClientCapabilities;
use mcp_types::Implementation;
use mcp_types::McpClientInfo;
use mcp_types::Tool;

use serde_json::json;
Expand Down Expand Up @@ -159,7 +159,7 @@ impl McpConnectionManager {
// indicates this should be an empty object.
elicitation: Some(json!({})),
},
client_info: Implementation {
client_info: McpClientInfo {
name: "codex-mcp-client".to_owned(),
version: env!("CARGO_PKG_VERSION").to_owned(),
title: Some("Codex".into()),
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/mcp-client/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,10 @@ use anyhow::Context;
use anyhow::Result;
use codex_mcp_client::McpClient;
use mcp_types::ClientCapabilities;
use mcp_types::Implementation;
use mcp_types::InitializeRequestParams;
use mcp_types::ListToolsRequestParams;
use mcp_types::MCP_SCHEMA_VERSION;
use mcp_types::McpClientInfo;
use tracing_subscriber::EnvFilter;

#[tokio::main]
Expand Down Expand Up @@ -60,7 +60,7 @@ async fn main() -> Result<()> {
sampling: None,
elicitation: None,
},
client_info: Implementation {
client_info: McpClientInfo {
name: "codex-mcp-client".to_owned(),
version: env!("CARGO_PKG_VERSION").to_owned(),
title: Some("Codex".to_string()),
Expand Down
1 change: 1 addition & 0 deletions codex-rs/mcp-server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ uuid = { version = "1", features = ["serde", "v4"] }
[dev-dependencies]
assert_cmd = "2"
mcp_test_support = { path = "tests/common" }
os_info = "3.12.0"
pretty_assertions = "1.4.1"
tempfile = "3"
wiremock = "0.6"
13 changes: 12 additions & 1 deletion codex-rs/mcp-server/src/message_processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ use codex_protocol::mcp_protocol::ConversationId;
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;
Expand Down Expand Up @@ -208,6 +210,14 @@ impl MessageProcessor {
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.
Expand All @@ -224,10 +234,11 @@ impl MessageProcessor {
},
instructions: None,
protocol_version: params.protocol_version.clone(),
server_info: mcp_types::Implementation {
server_info: mcp_types::McpServerInfo {
name: "codex-mcp-server".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
title: Some("Codex".to_string()),
user_agent: get_codex_user_agent(),
},
};

Expand Down
2 changes: 2 additions & 0 deletions codex-rs/mcp-server/tests/common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,11 @@ path = "lib.rs"
[dependencies]
anyhow = "1"
assert_cmd = "2"
codex-core = { path = "../../../core" }
codex-mcp-server = { path = "../.." }
codex-protocol = { path = "../../../protocol" }
mcp-types = { path = "../../../mcp-types" }
os_info = "3.12.0"
pretty_assertions = "1.4.1"
serde = { version = "1" }
serde_json = "1"
Expand Down
15 changes: 12 additions & 3 deletions codex-rs/mcp-server/tests/common/mcp_process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,13 @@ use codex_protocol::mcp_protocol::SendUserTurnParams;

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::McpClientInfo;
use mcp_types::ModelContextProtocolNotification;
use mcp_types::ModelContextProtocolRequest;
use mcp_types::RequestId;
Expand Down Expand Up @@ -111,7 +111,7 @@ impl McpProcess {
roots: None,
sampling: None,
},
client_info: Implementation {
client_info: McpClientInfo {
name: "elicitation test".into(),
title: Some("Elicitation Test".into()),
version: "0.0.0".into(),
Expand All @@ -129,6 +129,14 @@ impl McpProcess {
.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(),
Expand All @@ -142,7 +150,8 @@ impl McpProcess {
"serverInfo": {
"name": "codex-mcp-server",
"title": "Codex",
"version": "0.0.0"
"version": "0.0.0",
"user_agent": user_agent
},
"protocolVersion": mcp_types::MCP_SCHEMA_VERSION
})
Expand Down
14 changes: 10 additions & 4 deletions codex-rs/mcp-server/tests/suite/user_agent.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use codex_core::default_client::get_codex_user_agent;
use codex_protocol::mcp_protocol::GetUserAgentResponse;
use mcp_test_support::McpProcess;
use mcp_test_support::to_response;
Expand Down Expand Up @@ -34,11 +33,18 @@ async fn get_user_agent_returns_current_codex_user_agent() {
.expect("getUserAgent timeout")
.expect("getUserAgent response");

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()
);

let received: GetUserAgentResponse =
to_response(response).expect("deserialize getUserAgent response");
let expected = GetUserAgentResponse {
user_agent: get_codex_user_agent(),
};
let expected = GetUserAgentResponse { user_agent };

assert_eq!(received, expected);
}
16 changes: 13 additions & 3 deletions codex-rs/mcp-types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -482,13 +482,23 @@ pub struct ImageContent {

/// Describes the name and version of an MCP implementation, with an optional title for UI representation.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, TS)]
pub struct Implementation {
pub struct McpClientInfo {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
pub version: String,
}

/// Describes the name and version of an MCP implementation, with an optional title for UI representation.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, TS)]
pub struct McpServerInfo {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
pub version: String,
pub user_agent: String,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, TS)]
pub enum InitializeRequest {}

Expand All @@ -502,7 +512,7 @@ impl ModelContextProtocolRequest for InitializeRequest {
pub struct InitializeRequestParams {
pub capabilities: ClientCapabilities,
#[serde(rename = "clientInfo")]
pub client_info: Implementation,
pub client_info: McpClientInfo,
#[serde(rename = "protocolVersion")]
pub protocol_version: String,
}
Expand All @@ -516,7 +526,7 @@ pub struct InitializeResult {
#[serde(rename = "protocolVersion")]
pub protocol_version: String,
#[serde(rename = "serverInfo")]
pub server_info: Implementation,
pub server_info: McpServerInfo,
}

impl From<InitializeResult> for serde_json::Value {
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/mcp-types/tests/suite/initialize.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use mcp_types::ClientCapabilities;
use mcp_types::ClientRequest;
use mcp_types::Implementation;
use mcp_types::InitializeRequestParams;
use mcp_types::JSONRPC_VERSION;
use mcp_types::JSONRPCMessage;
use mcp_types::JSONRPCRequest;
use mcp_types::McpClientInfo;
use mcp_types::RequestId;
use serde_json::json;

Expand Down Expand Up @@ -58,7 +58,7 @@ fn deserialize_initialize_request() {
sampling: None,
elicitation: None,
},
client_info: Implementation {
client_info: McpClientInfo {
name: "acme-client".into(),
title: Some("Acme".to_string()),
version: "1.2.3".into(),
Expand Down
1 change: 1 addition & 0 deletions codex-rs/protocol-ts/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ path = "src/main.rs"

[dependencies]
anyhow = "1"
mcp-types = { path = "../mcp-types" }
codex-protocol = { path = "../protocol" }
ts-rs = "11"
clap = { version = "4", features = ["derive"] }
1 change: 1 addition & 0 deletions codex-rs/protocol-ts/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ pub fn generate_ts(out_dir: &Path, prettier: Option<&Path>) -> Result<()> {
ensure_dir(out_dir)?;

// Generate TS bindings
mcp_types::InitializeResult::export_all_to(out_dir)?;
codex_protocol::mcp_protocol::ConversationId::export_all_to(out_dir)?;
codex_protocol::mcp_protocol::InputItem::export_all_to(out_dir)?;
codex_protocol::mcp_protocol::ClientRequest::export_all_to(out_dir)?;
Expand Down
Loading