Increase platform abstraction cohesion
This commit is contained in:
387
rust/manticore-mcp-worker/src/main.rs
Normal file
387
rust/manticore-mcp-worker/src/main.rs
Normal file
@@ -0,0 +1,387 @@
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::env;
|
||||
use std::io::{self, BufRead, Read, Write};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn now_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn sql_escape(value: &str) -> String {
|
||||
value.replace('\\', "\\\\").replace('\'', "\\'")
|
||||
}
|
||||
|
||||
fn to_string_vec(value: &Value) -> Vec<String> {
|
||||
match value {
|
||||
Value::Array(items) => items
|
||||
.iter()
|
||||
.filter_map(|item| item.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect(),
|
||||
Value::String(raw) => raw
|
||||
.split(',')
|
||||
.map(|part| part.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect(),
|
||||
_ => vec![],
|
||||
}
|
||||
}
|
||||
|
||||
struct ManticoreClient {
|
||||
base_url: String,
|
||||
table: String,
|
||||
timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl ManticoreClient {
|
||||
fn from_env() -> Self {
|
||||
let base_url = env::var("MANTICORE_HTTP_URL")
|
||||
.unwrap_or_else(|_| "http://127.0.0.1:9308".to_string())
|
||||
.trim_end_matches('/')
|
||||
.to_string();
|
||||
let table = env::var("MANTICORE_MEMORY_TABLE")
|
||||
.unwrap_or_else(|_| "gia_memory_items".to_string())
|
||||
.trim()
|
||||
.to_string();
|
||||
let timeout_secs = env::var("MANTICORE_HTTP_TIMEOUT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.unwrap_or(5);
|
||||
Self {
|
||||
base_url,
|
||||
table,
|
||||
timeout_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn sql(&self, query: &str) -> Result<Value, String> {
|
||||
let endpoint = format!("{}/sql", self.base_url);
|
||||
let response = ureq::post(&endpoint)
|
||||
.timeout(std::time::Duration::from_secs(self.timeout_secs))
|
||||
.set("Content-Type", "application/x-www-form-urlencoded")
|
||||
.send_string(&format!(
|
||||
"mode=raw&query={}",
|
||||
urlencoding::encode(query).into_owned()
|
||||
))
|
||||
.map_err(|err| err.to_string())?;
|
||||
let body = response
|
||||
.into_string()
|
||||
.map_err(|err| format!("manticore response read failed: {err}"))?;
|
||||
serde_json::from_str::<Value>(&body)
|
||||
.map_err(|err| format!("manticore response parse failed: {err}"))
|
||||
}
|
||||
|
||||
fn ensure_table(&self) -> Result<(), String> {
|
||||
let query = format!(
|
||||
"CREATE TABLE IF NOT EXISTS {} (id BIGINT,memory_uuid STRING,user_id BIGINT,conversation_id STRING,memory_kind STRING,status STRING,updated_ts BIGINT,summary TEXT,body TEXT)",
|
||||
self.table
|
||||
);
|
||||
self.sql(&query).map(|_| ())
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_specs() -> Value {
|
||||
json!([
|
||||
{
|
||||
"name": "manticore.status",
|
||||
"description": "Report Manticore connectivity and table status.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"properties": {}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "manticore.query",
|
||||
"description": "Run fast full-text retrieval against the Manticore memory table.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string"},
|
||||
"user_id": {"type": "integer"},
|
||||
"conversation_id": {"type": "string"},
|
||||
"statuses": {
|
||||
"oneOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}}
|
||||
]
|
||||
},
|
||||
"limit": {"type": "integer", "minimum": 1, "maximum": 100}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "manticore.reindex",
|
||||
"description": "Run table maintenance operations for fast reads (flush + optimize).",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"flush_ramchunk": {"type": "boolean"},
|
||||
"optimize": {"type": "boolean"}
|
||||
}
|
||||
}
|
||||
}
|
||||
])
|
||||
}
|
||||
|
||||
fn call_tool(client: &ManticoreClient, name: &str, arguments: &Value) -> Result<Value, String> {
|
||||
match name {
|
||||
"manticore.status" => {
|
||||
client.ensure_table()?;
|
||||
let table_check = client.sql(&format!(
|
||||
"SHOW TABLES LIKE '{}'",
|
||||
sql_escape(&client.table)
|
||||
))?;
|
||||
Ok(json!({
|
||||
"backend": "manticore",
|
||||
"ok": true,
|
||||
"manticore_http_url": client.base_url,
|
||||
"manticore_table": client.table,
|
||||
"table_check": table_check,
|
||||
"ts_ms": now_ms()
|
||||
}))
|
||||
}
|
||||
"manticore.query" => {
|
||||
client.ensure_table()?;
|
||||
let query = arguments
|
||||
.get("query")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
if query.is_empty() {
|
||||
return Err("query is required".to_string());
|
||||
}
|
||||
let limit = arguments
|
||||
.get("limit")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(20)
|
||||
.clamp(1, 100);
|
||||
let mut where_parts = vec![format!("MATCH('{}')", sql_escape(&query))];
|
||||
if let Some(user_id) = arguments.get("user_id").and_then(Value::as_i64) {
|
||||
where_parts.push(format!("user_id={user_id}"));
|
||||
}
|
||||
let conversation_id = arguments
|
||||
.get("conversation_id")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
if !conversation_id.is_empty() {
|
||||
where_parts.push(format!(
|
||||
"conversation_id='{}'",
|
||||
sql_escape(&conversation_id)
|
||||
));
|
||||
}
|
||||
let statuses = to_string_vec(arguments.get("statuses").unwrap_or(&Value::Null));
|
||||
if !statuses.is_empty() {
|
||||
let joined = statuses
|
||||
.iter()
|
||||
.map(|s| format!("'{}'", sql_escape(s)))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
where_parts.push(format!("status IN ({joined})"));
|
||||
}
|
||||
let sql = format!(
|
||||
"SELECT memory_uuid,memory_kind,status,conversation_id,updated_ts,summary,WEIGHT() AS score FROM {} WHERE {} ORDER BY score DESC LIMIT {}",
|
||||
client.table,
|
||||
where_parts.join(" AND "),
|
||||
limit
|
||||
);
|
||||
let payload = client.sql(&sql)?;
|
||||
let count = payload
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.map(|rows| rows.len())
|
||||
.unwrap_or(0);
|
||||
Ok(json!({
|
||||
"backend": "manticore",
|
||||
"query": query,
|
||||
"count": count,
|
||||
"hits": payload.get("data").cloned().unwrap_or_else(|| json!([])),
|
||||
"raw": payload
|
||||
}))
|
||||
}
|
||||
"manticore.reindex" => {
|
||||
client.ensure_table()?;
|
||||
let flush = arguments
|
||||
.get("flush_ramchunk")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
let optimize = arguments
|
||||
.get("optimize")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
let mut actions: Vec<Value> = Vec::new();
|
||||
if flush {
|
||||
let sql = format!("FLUSH RAMCHUNK {}", client.table);
|
||||
let payload = client.sql(&sql)?;
|
||||
actions.push(json!({"sql": sql, "result": payload}));
|
||||
}
|
||||
if optimize {
|
||||
let sql = format!("OPTIMIZE TABLE {}", client.table);
|
||||
let payload = client.sql(&sql)?;
|
||||
actions.push(json!({"sql": sql, "result": payload}));
|
||||
}
|
||||
Ok(json!({
|
||||
"ok": true,
|
||||
"actions": actions,
|
||||
"ts_ms": now_ms()
|
||||
}))
|
||||
}
|
||||
_ => Err(format!("Unknown tool: {name}")),
|
||||
}
|
||||
}
|
||||
|
||||
fn response(id: Value, result: Value) -> Value {
|
||||
json!({"jsonrpc":"2.0","id":id,"result":result})
|
||||
}
|
||||
|
||||
fn error(id: Value, code: i32, message: &str) -> Value {
|
||||
json!({"jsonrpc":"2.0","id":id,"error":{"code":code,"message":message}})
|
||||
}
|
||||
|
||||
fn write_message(payload: &Value, compat_newline_mode: bool) -> Result<(), String> {
|
||||
let raw = serde_json::to_vec(payload).map_err(|e| e.to_string())?;
|
||||
let mut stdout = io::stdout();
|
||||
if compat_newline_mode {
|
||||
stdout
|
||||
.write_all(format!("{}\n", String::from_utf8_lossy(&raw)).as_bytes())
|
||||
.map_err(|e| e.to_string())?;
|
||||
} else {
|
||||
stdout
|
||||
.write_all(format!("Content-Length: {}\r\n\r\n", raw.len()).as_bytes())
|
||||
.map_err(|e| e.to_string())?;
|
||||
stdout.write_all(&raw).map_err(|e| e.to_string())?;
|
||||
}
|
||||
stdout.flush().map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
fn read_message(
|
||||
stdin: &mut io::StdinLock<'_>,
|
||||
compat_newline_mode: &mut bool,
|
||||
) -> Result<Option<Value>, String> {
|
||||
let mut headers: HashMap<String, String> = HashMap::new();
|
||||
let mut pending_body: Vec<u8> = Vec::new();
|
||||
loop {
|
||||
let mut line: Vec<u8> = Vec::new();
|
||||
let bytes = stdin.read_until(b'\n', &mut line).map_err(|e| e.to_string())?;
|
||||
if bytes == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
let trimmed = line
|
||||
.iter()
|
||||
.copied()
|
||||
.skip_while(|b| b.is_ascii_whitespace())
|
||||
.collect::<Vec<u8>>();
|
||||
if headers.is_empty() && (trimmed.starts_with(b"{") || trimmed.starts_with(b"[")) {
|
||||
*compat_newline_mode = true;
|
||||
let raw = String::from_utf8_lossy(&line).trim().to_string();
|
||||
let parsed = serde_json::from_str::<Value>(&raw).map_err(|e| e.to_string())?;
|
||||
return Ok(Some(parsed));
|
||||
}
|
||||
|
||||
if line == b"\r\n" || line == b"\n" {
|
||||
break;
|
||||
}
|
||||
let decoded = String::from_utf8_lossy(&line).trim().to_string();
|
||||
if let Some((k, v)) = decoded.split_once(':') {
|
||||
headers.insert(k.trim().to_lowercase(), v.trim().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(length_raw) = headers.get("content-length") {
|
||||
let length = length_raw
|
||||
.parse::<usize>()
|
||||
.map_err(|_| "invalid content-length".to_string())?;
|
||||
if length > 0 {
|
||||
pending_body.resize(length, 0);
|
||||
stdin.read_exact(&mut pending_body).map_err(|e| e.to_string())?;
|
||||
let parsed =
|
||||
serde_json::from_slice::<Value>(&pending_body).map_err(|e| e.to_string())?;
|
||||
return Ok(Some(parsed));
|
||||
}
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let client = ManticoreClient::from_env();
|
||||
let stdin = io::stdin();
|
||||
let mut locked = stdin.lock();
|
||||
let mut compat_newline_mode = false;
|
||||
|
||||
loop {
|
||||
let message = match read_message(&mut locked, &mut compat_newline_mode) {
|
||||
Ok(Some(value)) => value,
|
||||
Ok(None) => return,
|
||||
Err(err) => {
|
||||
let _ = write_message(
|
||||
&error(Value::Null, -32000, &format!("read failed: {err}")),
|
||||
compat_newline_mode,
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let id = message.get("id").cloned().unwrap_or(Value::Null);
|
||||
let method = message
|
||||
.get("method")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let params = message.get("params").cloned().unwrap_or_else(|| json!({}));
|
||||
|
||||
let response_payload = match method.as_str() {
|
||||
"notifications/initialized" => None,
|
||||
"initialize" => Some(response(
|
||||
id,
|
||||
json!({
|
||||
"protocolVersion":"2025-06-18",
|
||||
"serverInfo":{"name":"gia-manticore-mcp-rust","version":"0.1.0"},
|
||||
"capabilities":{"tools":{}}
|
||||
}),
|
||||
)),
|
||||
"ping" => Some(response(id, json!({}))),
|
||||
"tools/list" => Some(response(id, json!({"tools": tool_specs()}))),
|
||||
"tools/call" => {
|
||||
let name = params
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
let args = params
|
||||
.get("arguments")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!({}));
|
||||
match call_tool(&client, &name, &args) {
|
||||
Ok(payload) => Some(response(
|
||||
id,
|
||||
json!({"isError":false,"content":[{"type":"text","text":payload.to_string()}]}),
|
||||
)),
|
||||
Err(err) => Some(response(
|
||||
id,
|
||||
json!({"isError":true,"content":[{"type":"text","text":json!({"error":err}).to_string()}]}),
|
||||
)),
|
||||
}
|
||||
}
|
||||
_ => Some(error(id, -32601, &format!("Method not found: {method}"))),
|
||||
};
|
||||
|
||||
if let Some(payload) = response_payload {
|
||||
if write_message(&payload, compat_newline_mode).is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user