mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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
This commit is contained in:
jif-oai
committed by
GitHub
Unverified
GitHub
parent
161541310f
commit
f20f8a719e
Generated
+1
@@ -3047,6 +3047,7 @@ dependencies = [
|
||||
"codex-utils-output-truncation",
|
||||
"pretty_assertions",
|
||||
"rmcp",
|
||||
"schemars 0.8.22",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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<String>,
|
||||
pub entries: Vec<MemoryEntry>,
|
||||
@@ -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<String>,
|
||||
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,
|
||||
|
||||
@@ -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<T: JsonSchema>() -> JsonObject {
|
||||
schema_for::<T>(/*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<T: JsonSchema>() -> JsonObject {
|
||||
schema_for::<T>(/*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<T: JsonSchema>(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::<T>();
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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<B> {
|
||||
tools: Arc<Vec<Tool>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[derive(Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ListArgs {
|
||||
path: Option<String>,
|
||||
cursor: Option<String>,
|
||||
#[schemars(range(min = 1))]
|
||||
max_results: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[derive(Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ReadArgs {
|
||||
path: String,
|
||||
#[schemars(range(min = 1))]
|
||||
line_offset: Option<usize>,
|
||||
#[schemars(range(min = 1))]
|
||||
max_lines: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct SearchArgs {
|
||||
#[schemars(length(min = 1))]
|
||||
queries: Vec<String>,
|
||||
match_mode: Option<SearchMatchMode>,
|
||||
path: Option<String>,
|
||||
cursor: Option<String>,
|
||||
#[schemars(range(min = 0))]
|
||||
context_lines: Option<usize>,
|
||||
case_sensitive: Option<bool>,
|
||||
#[schemars(range(min = 1))]
|
||||
max_results: Option<usize>,
|
||||
}
|
||||
|
||||
@@ -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::<ListArgs>()),
|
||||
);
|
||||
tool.output_schema = Some(Arc::new(schema::list_output_schema()));
|
||||
tool.output_schema = Some(Arc::new(schema::output_schema_for::<ListMemoriesResponse>()));
|
||||
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::<ReadArgs>()),
|
||||
);
|
||||
tool.output_schema = Some(Arc::new(schema::read_output_schema()));
|
||||
tool.output_schema = Some(Arc::new(schema::output_schema_for::<ReadMemoryResponse>()));
|
||||
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::<SearchArgs>()),
|
||||
);
|
||||
tool.output_schema = Some(Arc::new(schema::search_output_schema()));
|
||||
tool.output_schema = Some(Arc::new(
|
||||
schema::output_schema_for::<SearchMemoriesResponse>(),
|
||||
));
|
||||
tool.annotations = Some(ToolAnnotations::new().read_only(true));
|
||||
tool
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user