use std::num::NonZeroU32;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use super::read_support::{Snapshot, bounded};
use super::{EncodeError, Sender, ServiceRefusal, ToolResult, refused};
use crate::mcp::{rpc, types};
use crate::serve::protocol;
#[derive(Deserialize)]
struct Collections {
selection: Snapshot,
entries: Vec<Collection>,
next_after: Option<String>,
}
#[derive(Deserialize)]
struct Collection {
name: String,
schema_json: Option<String>,
}
#[derive(Deserialize)]
struct Branches {
entries: Vec<Snapshot>,
next_after: Option<String>,
}
#[derive(Deserialize)]
struct Relations {
selection: Snapshot,
relations: Vec<types::RelationInfo>,
}
#[derive(Deserialize)]
struct Vectors {
selection: Snapshot,
fields: Vec<Vector>,
embedded_fields: Vec<Vector>,
}
#[derive(Deserialize)]
struct Vector {
collection: String,
field: String,
dimensions: std::num::NonZeroU64,
#[serde(default)]
count: Option<u64>,
}
#[derive(Deserialize)]
struct Counts {
selection: Snapshot,
collection_count: Option<u64>,
documents: Option<u64>,
edges: Option<u64>,
vectors: Option<u64>,
}
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Cursor {
branch: String,
version: types::Version,
part: String,
after: String,
}
fn refusal(code: &str, message: impl Into<String>) -> ServiceRefusal {
ServiceRefusal {
code: code.to_owned(),
message: message.into(),
receipt: None,
}
}
fn answer<T: DeserializeOwned>(
sender: &mut impl Sender,
request: protocol::Operation,
) -> Result<T, ServiceRefusal> {
let raw = sender.send(request)?;
serde_json::from_str(raw.get())
.map_err(|error| refusal("mcp_service_answer_invalid", error.to_string()))
}
fn cursor(
value: Option<&str>,
selected: Option<&Snapshot>,
part: &str,
) -> Result<Option<String>, ServiceRefusal> {
let Some(value) = value else { return Ok(None) };
let selected = selected.ok_or_else(|| {
refusal(
"mcp_version_required",
"Continuation requires the captured selection version",
)
})?;
let cursor: Cursor = serde_json::from_str(value).map_err(|_| {
refusal(
"mcp_cursor_invalid",
"The cursor is not a typed store-description cursor",
)
})?;
if cursor.branch != selected.branch
|| cursor.version.0 != selected.version.0
|| cursor.part != part
{
return Err(refusal(
"mcp_cursor_invalid",
"The cursor does not identify this captured selection and page part",
));
}
Ok(Some(cursor.after))
}
fn token(
after: Option<String>,
selected: &Snapshot,
part: &str,
) -> Result<Option<String>, ServiceRefusal> {
after
.map(|after| {
serde_json::to_string(&Cursor {
branch: selected.branch.clone(),
version: selected.version,
part: part.to_owned(),
after,
})
.map_err(|error| refusal("mcp_cursor_encoding_failed", error.to_string()))
})
.transpose()
}
fn page<T>(
mut entries: Vec<T>,
after: Option<&str>,
limit: usize,
key: impl Fn(&T) -> String,
) -> (Vec<T>, Option<String>) {
entries.sort_by_key(&key);
let mut selected = entries
.into_iter()
.filter(|entry| after.is_none_or(|after| key(entry).as_str() > after));
let entries: Vec<_> = selected.by_ref().take(limit).collect();
let next = selected.next().and_then(|_| entries.last().map(&key));
(entries, next)
}
pub(super) fn dispatch(
sender: &mut impl Sender,
call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
let input: types::DescribeStoreInput = match serde_json::from_str(call.arguments.get()) {
Ok(input) => input,
Err(error) => return refused(&refusal("mcp_arguments_invalid", error.to_string())),
};
match describe(sender, &input) {
Ok(output) => Ok(output),
Err(refusal) => refused(&refusal),
}
}
struct Description {
selected: Snapshot,
collections: Vec<Collection>,
next_collection: Option<String>,
branches: Branches,
relations: Vec<types::RelationInfo>,
vectors: Vec<Vector>,
counts: Counts,
}
fn limit(input: &types::DescribeStoreInput) -> Result<NonZeroU32, ServiceRefusal> {
let has_cursor = input.after.branches.is_some()
|| input.after.collections.is_some()
|| input.after.relations.is_some()
|| input.after.vector_fields.is_some();
if has_cursor && input.selection.version.is_none() {
return Err(refusal(
"mcp_version_required",
"Continuation requires the captured selection version",
));
}
Ok(input.page.rows)
}
fn read(
sender: &mut impl Sender,
selection: &types::Selection,
cursors: &types::DescribeCursors,
limit: NonZeroU32,
) -> Result<Description, ServiceRefusal> {
let expected = selection.version.map(|version| Snapshot {
branch: selection.branch.clone(),
version,
});
let collections: Collections = answer(
sender,
protocol::Operation::Collections(protocol::CollectionPage {
selection: protocol::Selection {
branch: selection.branch.clone(),
version: selection.version.map(|version| version.0),
},
after: cursor(
cursors.collections.as_deref(),
expected.as_ref(),
"collections",
)?,
limit,
}),
)?;
collections.selection.requested(selection)?;
let selected = collections.selection;
let branches: Branches = answer(
sender,
protocol::Operation::Branches(protocol::BranchPage {
after: cursor(cursors.branches.as_deref(), Some(&selected), "branches")?,
limit,
}),
)?;
let relations: Relations = answer(
sender,
protocol::Operation::Relations(protocol::RelationsRead {
selection: selected.selection(),
}),
)?;
relations.selection.require(&selected)?;
let vectors: Vectors = answer(
sender,
protocol::Operation::VectorFields(protocol::StoreRead {
selection: selected.selection(),
}),
)?;
vectors.selection.require(&selected)?;
let counts: Counts = answer(
sender,
protocol::Operation::StoreStats(protocol::StoreRead {
selection: selected.selection(),
}),
)?;
counts.selection.require(&selected)?;
Ok(Description {
selected,
collections: collections.entries,
next_collection: collections.next_after,
branches,
relations: relations.relations,
vectors: vectors
.fields
.into_iter()
.chain(vectors.embedded_fields)
.collect(),
counts,
})
}
fn counts(counts: &Counts) -> Result<types::Counts, ServiceRefusal> {
let known = |value: Option<u64>| {
value.ok_or_else(|| {
refusal(
"mcp_counts_unavailable",
"Kept store counts are unavailable; no count has been fabricated",
)
})
};
Ok(types::Counts {
collections: known(counts.collection_count)?,
documents: known(counts.documents)?,
edges: known(counts.edges)?,
vectors: known(counts.vectors)?,
})
}
fn collection(entry: Collection) -> Result<types::CollectionInfo, ServiceRefusal> {
let schema = match entry.schema_json {
Some(raw) => serde_json::from_str(&raw)
.map_err(|error| refusal("mcp_service_answer_invalid", error.to_string()))?,
None => types::CollectionShape::Remove,
};
Ok(types::CollectionInfo {
name: entry.name,
schema,
})
}
fn project(
description: Description,
cursors: &types::DescribeCursors,
limit: NonZeroU32,
) -> Result<types::DescribeStoreOutput, ServiceRefusal> {
let bound = usize::try_from(limit.get()).map_err(|_| {
refusal(
"mcp_page_rows_unrepresentable",
"The requested rows do not fit this address space",
)
})?;
let selected = description.selected;
let (relations, next_relation) = page(
description.relations,
cursor(cursors.relations.as_deref(), Some(&selected), "relations")?.as_deref(),
bound,
|entry| entry.relation.clone(),
);
let vectors: Vec<_> = description
.vectors
.into_iter()
.map(|entry| types::VectorField {
collection: entry.collection,
field: entry.field,
dimensions: entry.dimensions,
count: entry.count,
})
.collect();
let (vectors, next_vector) = page(
vectors,
cursor(
cursors.vector_fields.as_deref(),
Some(&selected),
"vector_fields",
)?
.as_deref(),
bound,
|entry| format!("{}\0{}", entry.collection, entry.field),
);
let next_branch = token(description.branches.next_after, &selected, "branches")?;
let next_collection = token(description.next_collection, &selected, "collections")?;
let next_relation = token(next_relation, &selected, "relations")?;
let next_vector = token(next_vector, &selected, "vector_fields")?;
Ok(types::DescribeStoreOutput {
selection: selected.output(),
branches: types::Page {
entries: description
.branches
.entries
.into_iter()
.map(Snapshot::output)
.collect(),
next_after: next_branch,
},
collections: types::Page {
entries: description
.collections
.into_iter()
.map(collection)
.collect::<Result<Vec<_>, _>>()?,
next_after: next_collection,
},
relations: types::Page {
entries: relations,
next_after: next_relation,
},
vector_fields: types::Page {
entries: vectors,
next_after: next_vector,
},
counts: counts(&description.counts)?,
})
}
fn describe(
sender: &mut impl Sender,
input: &types::DescribeStoreInput,
) -> Result<ToolResult, ServiceRefusal> {
let limit = limit(input)?;
let description = read(sender, &input.selection, &input.after, limit)?;
let output = project(description, &input.after, limit)?;
bounded(&output, input.page.bytes.get())
}