From f20f8a719e8ab8ede57754d9ded7ba4ba28f8b4f Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 4 May 2026 18:40:17 +0200 Subject: [PATCH] memories/mcp: generate tool schemas with schemars (#21012) ## Why The memories MCP server currently keeps handwritten JSON Schema beside the Rust types that actually serialize and deserialize the tool payloads: [`schema.rs`](https://github.com/openai/codex/blob/2f5c06a29cdd68f11d07126dc56871bff1218ba1/codex-rs/memories/mcp/src/schema.rs#L4-L133), [`server.rs`](https://github.com/openai/codex/blob/2f5c06a29cdd68f11d07126dc56871bff1218ba1/codex-rs/memories/mcp/src/server.rs#L44-L75), and [`backend.rs`](https://github.com/openai/codex/blob/2f5c06a29cdd68f11d07126dc56871bff1218ba1/codex-rs/memories/mcp/src/backend.rs#L41-L117). That duplicates the tool contract and makes schema drift easier as the API evolves. ## What changed - derive `JsonSchema` for the memories tool arguments, responses, and nested response types - replace the handwritten schema builders with shared `schemars` generation - preserve the existing wire shape while generating schemas, including nullable output `Option` fields and non-nullable optional input fields - wire the `list`, `read`, and `search` tools to the generated schemas ## Verification - CI pending --- codex-rs/Cargo.lock | 1 + codex-rs/memories/mcp/Cargo.toml | 1 + codex-rs/memories/mcp/src/backend.rs | 20 ++-- codex-rs/memories/mcp/src/schema.rs | 162 ++++++--------------------- codex-rs/memories/mcp/src/server.rs | 32 ++++-- 5 files changed, 72 insertions(+), 144 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c94785df6..53a1c4cc6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3047,6 +3047,7 @@ dependencies = [ "codex-utils-output-truncation", "pretty_assertions", "rmcp", + "schemars 0.8.22", "serde", "serde_json", "tempfile", diff --git a/codex-rs/memories/mcp/Cargo.toml b/codex-rs/memories/mcp/Cargo.toml index da1d46579..4e91472d4 100644 --- a/codex-rs/memories/mcp/Cargo.toml +++ b/codex-rs/memories/mcp/Cargo.toml @@ -19,6 +19,7 @@ rmcp = { workspace = true, default-features = false, features = [ "schemars", "server", ] } +schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/codex-rs/memories/mcp/src/backend.rs b/codex-rs/memories/mcp/src/backend.rs index 6cc3bd2cf..67e591c8d 100644 --- a/codex-rs/memories/mcp/src/backend.rs +++ b/codex-rs/memories/mcp/src/backend.rs @@ -1,3 +1,4 @@ +use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; use std::future::Future; @@ -38,7 +39,8 @@ pub struct ListMemoriesRequest { pub max_results: usize, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)] +#[schemars(deny_unknown_fields)] pub struct ListMemoriesResponse { pub path: Option, pub entries: Vec, @@ -54,7 +56,8 @@ pub struct ReadMemoryRequest { pub max_tokens: usize, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)] +#[schemars(deny_unknown_fields)] pub struct ReadMemoryResponse { pub path: String, pub start_line_number: usize, @@ -73,7 +76,8 @@ pub struct SearchMemoriesRequest { pub max_results: usize, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)] +#[schemars(deny_unknown_fields)] pub struct SearchMemoriesResponse { pub queries: Vec, pub match_mode: SearchMatchMode, @@ -83,27 +87,29 @@ pub struct SearchMemoriesResponse { pub truncated: bool, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum SearchMatchMode { Any, All, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)] +#[schemars(deny_unknown_fields)] pub struct MemoryEntry { pub path: String, pub entry_type: MemoryEntryType, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum MemoryEntryType { File, Directory, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)] +#[schemars(deny_unknown_fields)] pub struct MemorySearchMatch { pub path: String, pub match_line_number: usize, diff --git a/codex-rs/memories/mcp/src/schema.rs b/codex-rs/memories/mcp/src/schema.rs index 3e212cf29..2f01d2c95 100644 --- a/codex-rs/memories/mcp/src/schema.rs +++ b/codex-rs/memories/mcp/src/schema.rs @@ -1,136 +1,42 @@ use rmcp::model::JsonObject; -use serde_json::json; +use schemars::JsonSchema; +use schemars::r#gen::SchemaSettings; -pub(crate) fn list_input_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "path": { "type": "string" }, - "cursor": { "type": "string" }, - "max_results": { "type": "integer", "minimum": 1 } - }, - "additionalProperties": false - })) +pub(crate) fn input_schema_for() -> JsonObject { + schema_for::(/*option_add_null_type*/ false) } -pub(crate) fn list_output_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "path": { - "anyOf": [{ "type": "string" }, { "type": "null" }] - }, - "next_cursor": { - "anyOf": [{ "type": "string" }, { "type": "null" }] - }, - "entries": { - "type": "array", - "items": { - "type": "object", - "properties": { - "path": { "type": "string" }, - "entry_type": { "type": "string", "enum": ["file", "directory"] } - }, - "required": ["path", "entry_type"], - "additionalProperties": false - } - }, - "truncated": { "type": "boolean" } - }, - "required": ["path", "entries", "next_cursor", "truncated"], - "additionalProperties": false - })) +pub(crate) fn output_schema_for() -> JsonObject { + schema_for::(/*option_add_null_type*/ true) } -pub(crate) fn read_input_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "path": { "type": "string" }, - "line_offset": { "type": "integer", "minimum": 1 }, - "max_lines": { "type": "integer", "minimum": 1 } - }, - "required": ["path"], - "additionalProperties": false - })) -} +fn schema_for(option_add_null_type: bool) -> JsonObject { + let schema = SchemaSettings::draft2019_09() + .with(|settings| { + settings.inline_subschemas = true; + settings.option_add_null_type = option_add_null_type; + }) + .into_generator() + .into_root_schema_for::(); + let schema_value = serde_json::to_value(schema) + .unwrap_or_else(|err| panic!("generated tool schema should serialize: {err}")); + let serde_json::Value::Object(mut schema_object) = schema_value else { + unreachable!("root tool schema must be an object"); + }; -pub(crate) fn read_output_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "path": { "type": "string" }, - "start_line_number": { "type": "integer" }, - "content": { "type": "string" }, - "truncated": { "type": "boolean" } - }, - "required": ["path", "start_line_number", "content", "truncated"], - "additionalProperties": false - })) -} - -pub(crate) fn search_input_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "queries": { - "type": "array", - "items": { "type": "string" }, - "minItems": 1 - }, - "match_mode": { "type": "string", "enum": ["any", "all"] }, - "path": { "type": "string" }, - "cursor": { "type": "string" }, - "context_lines": { "type": "integer", "minimum": 0 }, - "case_sensitive": { "type": "boolean" }, - "max_results": { "type": "integer", "minimum": 1 } - }, - "required": ["queries"], - "additionalProperties": false - })) -} - -pub(crate) fn search_output_schema() -> JsonObject { - json_schema(json!({ - "type": "object", - "properties": { - "queries": { - "type": "array", - "items": { "type": "string" } - }, - "match_mode": { "type": "string", "enum": ["any", "all"] }, - "path": { - "anyOf": [{ "type": "string" }, { "type": "null" }] - }, - "next_cursor": { - "anyOf": [{ "type": "string" }, { "type": "null" }] - }, - "matches": { - "type": "array", - "items": { - "type": "object", - "properties": { - "path": { "type": "string" }, - "match_line_number": { "type": "integer" }, - "content_start_line_number": { "type": "integer" }, - "content": { "type": "string" }, - "matched_queries": { - "type": "array", - "items": { "type": "string" } - } - }, - "required": ["path", "match_line_number", "content_start_line_number", "content", "matched_queries"], - "additionalProperties": false - } - }, - "truncated": { "type": "boolean" } - }, - "required": ["queries", "match_mode", "path", "matches", "next_cursor", "truncated"], - "additionalProperties": false - })) -} - -fn json_schema(value: serde_json::Value) -> JsonObject { - serde_json::from_value(value) - .unwrap_or_else(|err| panic!("static tool schema should deserialize: {err}")) + // MCP tools only need the JSON Schema body, not schemars' root metadata. + let mut tool_schema = JsonObject::new(); + for key in [ + "properties", + "required", + "type", + "additionalProperties", + "$defs", + "definitions", + ] { + if let Some(value) = schema_object.remove(key) { + tool_schema.insert(key.to_string(), value); + } + } + tool_schema } diff --git a/codex-rs/memories/mcp/src/server.rs b/codex-rs/memories/mcp/src/server.rs index 2633d9448..e3a6db8fe 100644 --- a/codex-rs/memories/mcp/src/server.rs +++ b/codex-rs/memories/mcp/src/server.rs @@ -2,13 +2,16 @@ use crate::backend::DEFAULT_LIST_MAX_RESULTS; use crate::backend::DEFAULT_READ_MAX_TOKENS; use crate::backend::DEFAULT_SEARCH_MAX_RESULTS; use crate::backend::ListMemoriesRequest; +use crate::backend::ListMemoriesResponse; use crate::backend::MAX_LIST_RESULTS; use crate::backend::MAX_SEARCH_RESULTS; use crate::backend::MemoriesBackend; use crate::backend::MemoriesBackendError; use crate::backend::ReadMemoryRequest; +use crate::backend::ReadMemoryResponse; use crate::backend::SearchMatchMode; use crate::backend::SearchMemoriesRequest; +use crate::backend::SearchMemoriesResponse; use crate::local::LocalMemoriesBackend; use crate::schema; use anyhow::Context; @@ -25,6 +28,7 @@ use rmcp::model::ServerCapabilities; use rmcp::model::ServerInfo; use rmcp::model::Tool; use rmcp::model::ToolAnnotations; +use schemars::JsonSchema; use serde::Deserialize; use serde_json::json; use std::borrow::Cow; @@ -40,29 +44,37 @@ pub struct MemoriesMcpServer { tools: Arc>, } -#[derive(Deserialize)] +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] struct ListArgs { path: Option, cursor: Option, + #[schemars(range(min = 1))] max_results: Option, } -#[derive(Deserialize)] +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] struct ReadArgs { path: String, + #[schemars(range(min = 1))] line_offset: Option, + #[schemars(range(min = 1))] max_lines: Option, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct SearchArgs { + #[schemars(length(min = 1))] queries: Vec, match_mode: Option, path: Option, cursor: Option, + #[schemars(range(min = 0))] context_lines: Option, case_sensitive: Option, + #[schemars(range(min = 1))] max_results: Option, } @@ -191,9 +203,9 @@ fn list_tool() -> Tool { Cow::Borrowed( "List immediate files and directories under a path in the Codex memories store.", ), - Arc::new(schema::list_input_schema()), + Arc::new(schema::input_schema_for::()), ); - tool.output_schema = Some(Arc::new(schema::list_output_schema())); + tool.output_schema = Some(Arc::new(schema::output_schema_for::())); tool.annotations = Some(ToolAnnotations::new().read_only(true)); tool } @@ -204,9 +216,9 @@ fn read_tool() -> Tool { Cow::Borrowed( "Read a Codex memory file by relative path, optionally starting at a 1-indexed line offset and limiting the number of lines returned.", ), - Arc::new(schema::read_input_schema()), + Arc::new(schema::input_schema_for::()), ); - tool.output_schema = Some(Arc::new(schema::read_output_schema())); + tool.output_schema = Some(Arc::new(schema::output_schema_for::())); tool.annotations = Some(ToolAnnotations::new().read_only(true)); tool } @@ -217,9 +229,11 @@ fn search_tool() -> Tool { Cow::Borrowed( "Search Codex memory files for line-based substring matches, optionally requiring any or all query substrings on the same line.", ), - Arc::new(schema::search_input_schema()), + Arc::new(schema::input_schema_for::()), ); - tool.output_schema = Some(Arc::new(schema::search_output_schema())); + tool.output_schema = Some(Arc::new( + schema::output_schema_for::(), + )); tool.annotations = Some(ToolAnnotations::new().read_only(true)); tool }