use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::cache::{load_cached, save_cache};
use crate::coerce::to_kebab;
use crate::error::Result;
use crate::model::CommandDef;
pub const DESC_TRUNCATE: usize = 80;
pub const INDEX_VERSION: u32 = 4;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolIndexEntry {
pub name: String,
pub tool_name: String,
pub description: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactIndex {
pub v: u32,
pub names: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tool_names: Vec<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub tool_overrides: BTreeMap<u32, String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub descs: Vec<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub postings: BTreeMap<String, Vec<u32>>,
}
fn default_tool_name_from_kebab(kebab: &str) -> String {
kebab.replace('-', "_")
}
fn tokens_of(name: &str) -> impl Iterator<Item = &str> {
name.split('-').filter(|t| !t.is_empty())
}
pub fn rebuild_postings(index: &mut CompactIndex) {
let mut postings: BTreeMap<String, Vec<u32>> = BTreeMap::new();
for (i, name) in index.names.iter().enumerate() {
for tok in tokens_of(name) {
postings.entry(tok.to_string()).or_default().push(i as u32);
}
}
index.postings = postings;
}
impl CompactIndex {
pub fn len(&self) -> usize {
self.names.len()
}
pub fn is_empty(&self) -> bool {
self.names.is_empty()
}
pub fn has_descs(&self) -> bool {
!self.descs.is_empty()
}
pub fn tool_name_at(&self, i: usize) -> String {
if let Some(t) = self.tool_overrides.get(&(i as u32)) {
return t.clone();
}
if let Some(t) = self
.tool_names
.get(i)
.map(String::as_str)
.filter(|s| !s.is_empty())
{
return t.to_string();
}
default_tool_name_from_kebab(&self.names[i])
}
pub fn desc_at(&self, i: usize) -> &str {
self.descs.get(i).map(String::as_str).unwrap_or("")
}
pub fn entry_at(&self, i: usize) -> ToolIndexEntry {
ToolIndexEntry {
name: self.names[i].clone(),
tool_name: self.tool_name_at(i),
description: self.desc_at(i).to_string(),
}
}
pub fn to_entries(&self) -> Vec<ToolIndexEntry> {
(0..self.len()).map(|i| self.entry_at(i)).collect()
}
pub fn to_light_tools_json(&self, entries: &[ToolIndexEntry]) -> Value {
Value::Array(
entries
.iter()
.map(|e| {
serde_json::json!({
"name": e.tool_name,
"description": e.description,
})
})
.collect(),
)
}
}
pub fn index_cache_key(tools_cache_key: &str) -> String {
format!("{tools_cache_key}_index")
}
fn truncate_desc(s: &str, max: usize) -> String {
if s.chars().count() <= max {
return s.to_string();
}
let mut out: String = s.chars().take(max).collect();
if let Some(i) = out.rfind(' ') {
out.truncate(i);
}
out.push_str("...");
out
}
pub fn build_compact_index(tools: &[Value], with_descs: bool) -> CompactIndex {
let mut rows: Vec<(String, String, String)> = Vec::with_capacity(tools.len());
for tool in tools {
let tool_name = tool
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string();
let name = to_kebab(&tool_name);
let description = if with_descs {
truncate_desc(
tool.get("description")
.and_then(|v| v.as_str())
.unwrap_or(""),
DESC_TRUNCATE,
)
} else {
String::new()
};
rows.push((name, tool_name, description));
}
rows.sort_by(|a, b| a.0.cmp(&b.0));
let mut names = Vec::with_capacity(rows.len());
let mut tool_overrides: BTreeMap<u32, String> = BTreeMap::new();
let mut descs = Vec::new();
let mut postings: BTreeMap<String, Vec<u32>> = BTreeMap::new();
for (i, (name, tool_name, description)) in rows.into_iter().enumerate() {
let guessed = default_tool_name_from_kebab(&name);
if tool_name != guessed {
tool_overrides.insert(i as u32, tool_name);
}
for tok in tokens_of(&name) {
postings.entry(tok.to_string()).or_default().push(i as u32);
}
names.push(name);
if with_descs {
descs.push(description);
}
}
CompactIndex {
v: INDEX_VERSION,
names,
tool_names: Vec::new(),
tool_overrides,
descs,
postings,
}
}
pub fn build_index(tools: &[Value]) -> Vec<ToolIndexEntry> {
build_compact_index(tools, true).to_entries()
}
pub fn save_index(tools_cache_key: &str, tools: &[Value]) -> Result<()> {
let mut index = build_compact_index(tools, false);
index.postings.clear();
index.descs.clear();
save_cache(
&index_cache_key(tools_cache_key),
&serde_json::to_value(&index)?,
)
}
pub fn load_compact_index(tools_cache_key: &str, ttl: u64) -> Result<Option<CompactIndex>> {
let Some(v) = load_cached(&index_cache_key(tools_cache_key), ttl)? else {
return Ok(None);
};
let ver = v.get("v").and_then(|x| x.as_u64());
if matches!(ver, Some(2) | Some(3) | Some(4)) && v.get("names").is_some() {
let mut idx: CompactIndex = serde_json::from_value(v)?;
if idx.postings.is_empty() {
rebuild_postings(&mut idx);
}
return Ok(Some(idx));
}
if let Ok(entries) = serde_json::from_value::<Vec<ToolIndexEntry>>(v) {
return Ok(Some(legacy_entries_to_compact(entries)));
}
Ok(None)
}
fn legacy_entries_to_compact(mut entries: Vec<ToolIndexEntry>) -> CompactIndex {
entries.sort_by(|a, b| a.name.cmp(&b.name));
let tools: Vec<Value> = entries
.iter()
.map(|e| {
serde_json::json!({
"name": e.tool_name,
"description": e.description,
})
})
.collect();
build_compact_index(&tools, true)
}
pub fn load_index(tools_cache_key: &str, ttl: u64) -> Result<Option<Vec<ToolIndexEntry>>> {
Ok(load_compact_index(tools_cache_key, ttl)?.map(|c| c.to_entries()))
}
pub fn save_tools_and_index(tools_cache_key: &str, tools: &[Value]) -> Result<()> {
save_cache(tools_cache_key, &Value::Array(tools.to_vec()))?;
save_index(tools_cache_key, tools)?;
Ok(())
}
pub fn index_to_commands(entries: &[ToolIndexEntry]) -> Vec<CommandDef> {
entries
.iter()
.map(|e| CommandDef {
name: e.name.clone(),
description: e.description.clone(),
tool_name: Some(e.tool_name.clone()),
..Default::default()
})
.collect()
}
pub fn tools_to_light_commands(tools: &[Value]) -> Vec<CommandDef> {
index_to_commands(&build_index(tools))
}
pub fn search_compact(index: &CompactIndex, pattern: &str) -> Vec<ToolIndexEntry> {
let p = pattern.to_lowercase();
let ids = if !p.is_empty() && !p.contains('-') && !p.contains(' ') {
if let Some(list) = index.postings.get(&p) {
list.clone()
} else {
let mut acc: Vec<u32> = index
.postings
.iter()
.filter(|(k, _)| k.starts_with(&p) || k.contains(&p))
.flat_map(|(_, v)| v.iter().copied())
.collect();
acc.sort_unstable();
acc.dedup();
if acc.is_empty() {
return scan_names(index, &p);
}
acc
}
} else {
return scan_names(index, &p);
};
ids.into_iter()
.filter_map(|i| {
let i = i as usize;
(i < index.len()).then(|| index.entry_at(i))
})
.collect()
}
fn scan_names(index: &CompactIndex, p: &str) -> Vec<ToolIndexEntry> {
index
.names
.iter()
.enumerate()
.filter(|(i, n)| {
n.contains(p)
|| index.tool_name_at(*i).to_lowercase().contains(p)
|| index.desc_at(*i).to_lowercase().contains(p)
})
.map(|(i, _)| index.entry_at(i))
.collect()
}
pub fn search_index(entries: &[ToolIndexEntry], pattern: &str) -> Vec<ToolIndexEntry> {
let p = pattern.to_lowercase();
entries
.iter()
.filter(|e| {
e.name.to_lowercase().contains(&p)
|| e.tool_name.to_lowercase().contains(&p)
|| e.description.to_lowercase().contains(&p)
})
.cloned()
.collect()
}
pub fn find_exact(index: &CompactIndex, name: &str) -> Option<ToolIndexEntry> {
index
.names
.binary_search_by(|n| n.as_str().cmp(name))
.ok()
.map(|i| index.entry_at(i))
}
pub fn try_commands_from_index(
tools_cache_key: &str,
ttl: u64,
search: Option<&str>,
require_descs: bool,
) -> Result<Option<Vec<CommandDef>>> {
let Some(index) = load_compact_index(tools_cache_key, ttl)? else {
return Ok(None);
};
if require_descs && !index.has_descs() {
return Ok(None);
}
let entries = if let Some(pat) = search {
search_compact(&index, pat)
} else {
index.to_entries()
};
Ok(Some(index_to_commands(&entries)))
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn build_and_search_postings() {
let tools = vec![
json!({"name": "workers_list", "description": "List workers", "inputSchema": {"type":"object","properties":{"a":{"type":"string"}}}}),
json!({"name": "dns_records", "description": "DNS stuff", "inputSchema": {}}),
json!({"name": "workers_scripts_get", "description": "Get script", "inputSchema": {}}),
];
let idx = build_compact_index(&tools, false);
assert_eq!(idx.len(), 3);
assert!(idx.names.windows(2).all(|w| w[0] <= w[1]));
assert!(idx.postings.contains_key("workers"));
let hit = search_compact(&idx, "workers");
assert_eq!(hit.len(), 2);
assert!(find_exact(&idx, "dns-records").is_some());
}
#[test]
fn v4_disk_omits_postings_rebuild_on_load() {
let tools = vec![
json!({"name": "workers_list", "description": "List workers"}),
json!({"name": "dns_records", "description": "DNS"}),
];
let mut disk = build_compact_index(&tools, false);
disk.postings.clear();
disk.descs.clear();
let v = serde_json::to_value(&disk).unwrap();
assert!(
v.get("postings").is_none()
|| v.get("postings").unwrap().as_object().unwrap().is_empty()
);
let mut loaded: CompactIndex = serde_json::from_value(v).unwrap();
assert!(loaded.postings.is_empty());
rebuild_postings(&mut loaded);
assert!(loaded.postings.contains_key("workers"));
assert_eq!(search_compact(&loaded, "workers").len(), 1);
}
#[test]
fn compact_disk_smaller_than_naive_array() {
let tools: Vec<Value> = (0..50)
.map(|i| {
json!({
"name": format!("workers_scripts_op_{i}"),
"description": "x".repeat(200),
"inputSchema": {"type":"object","properties":{"a":{"type":"string"},"b":{"type":"integer"}}},
})
})
.collect();
let mut disk = build_compact_index(&tools, false);
disk.postings.clear();
let compact = serde_json::to_vec(&disk).unwrap();
let naive = serde_json::to_vec(&build_index(&tools)).unwrap();
assert!(
compact.len() < naive.len() / 2,
"disk compact {} vs naive {}",
compact.len(),
naive.len()
);
}
#[test]
fn snake_default_skips_tool_names_array() {
let tools = vec![
json!({"name": "workers_list", "description": "a"}),
json!({"name": "get_accounts_scim_v2_Groups", "description": "b"}),
];
let idx = build_compact_index(&tools, false);
assert!(idx.tool_names.is_empty());
let workers = idx.names.iter().position(|n| n == "workers-list").unwrap();
assert_eq!(idx.tool_name_at(workers), "workers_list");
let groups = idx
.names
.iter()
.position(|n| n == "get-accounts-scim-v2-groups")
.unwrap();
assert_eq!(idx.tool_name_at(groups), "get_accounts_scim_v2_Groups");
assert_eq!(idx.tool_overrides.len(), 1);
}
#[test]
fn require_descs_skips_names_only_index() {
let tools = vec![json!({"name": "workers_list", "description": "hi"})];
let idx = build_compact_index(&tools, false);
assert!(!idx.has_descs());
}
}