use std::sync::Arc;
use serde::Deserialize;
use serde_json::{json, Value};
use std::time::Instant;
use crate::compiler_pipeline;
use crate::compose;
use crate::knowledge::{Catalog, Category};
use crate::server::JsonRpcError;
use crate::telemetry::Telemetry;
pub fn list() -> Vec<Value> {
vec![
json!({
"name": "axon.primitives",
"description": "List every AXON primitive known to this server. \
Optionally filter by category. Use this FIRST when you need \
to know what's available before composing a program. \
Returns a sorted array of {name, summary, category, top_level, since}.",
"inputSchema": {
"type": "object",
"properties": {
"category": {
"type": "string",
"enum": [
"cognition", "cognitive_io", "data_plane",
"session_types", "wire", "operators"
],
"description": "Limit to one primitive family. Omit to list all."
}
},
"additionalProperties": false
}
}),
json!({
"name": "axon.primitive_doc",
"description": "Fetch the full reference for one AXON primitive by canonical \
name. Returns its grammar fragment, top-level status, the cycle that \
introduced it, and the complete markdown body (semantic constraints, \
examples, what-it-is-not, see-also). ALWAYS call this before generating \
a program that uses the primitive — the grammar is precise and the \
diagnostics are strict.",
"inputSchema": {
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Canonical primitive name (e.g. \"persona\", \
\"flow\", \"socket\", \"axonendpoint\")."
}
},
"required": ["name"],
"additionalProperties": false
}
}),
json!({
"name": "axon.check",
"description": "Validate AXON source code. Runs the same lex → parse → \
type-check pipeline the `axon check` CLI uses, and returns structured \
diagnostics (severity + stage + message + line/column). Use this AFTER \
generating a program, BEFORE presenting it to the user as working. \
Never claim a program is correct until axon.check returns ok=true. The \
pipeline stops at the first failing stage: a lex error means parse + \
type-check did not run.",
"inputSchema": {
"type": "object",
"properties": {
"source": {
"type": "string",
"description": "The complete AXON source text to validate. Pass \
the whole program (top-level declarations only — fragments \
are usually rejected)."
},
"filename": {
"type": "string",
"description": "Optional virtual filename for the diagnostics. \
Defaults to \"<axon.check input>\"."
}
},
"required": ["source"],
"additionalProperties": false
}
}),
json!({
"name": "axon.parse",
"description": "Parse AXON source to its Intermediate Representation (IR). \
Returns the same JSON shape `axon parse` would print: a `Program` node \
whose children are every declared primitive (`flows`, `personas`, \
`axonendpoints`, `axonstores`, `sockets`, `sessions`, …). Use this when \
you need to REASON ABOUT what you just wrote — e.g. to confirm the \
program has the right shape before refining it. On failure returns the \
same diagnostic shape as axon.check.",
"inputSchema": {
"type": "object",
"properties": {
"source": {
"type": "string",
"description": "The complete AXON source text to parse."
},
"filename": {
"type": "string",
"description": "Optional virtual filename. Defaults to \
\"<axon.parse input>\"."
}
},
"required": ["source"],
"additionalProperties": false
}
}),
json!({
"name": "axon.compose",
"description": "Generate a typed AXON scaffold from a natural-language intent. \
Classifies the intent into one of 8 closed domains (generic, healthcare, \
banking, government, legal, chat, retrieval, multi_agent), fetches the \
hand-authored template for that domain, re-validates it through the same \
`axon-frontend` pipeline `axon.check` uses, and returns: \
`{scaffold, domain, alternatives, primitives_used, compliance_applied, \
next_steps, axon_check_verdict}`. \
USE THIS when a user describes WHAT they want in plain language; the result \
is a guaranteed-compile starting point you can iterate on. The classifier \
is keyword-based and explainable — `alternatives` carries the scoreboard. \
Override the classifier with `domain:` when you already know the domain.",
"inputSchema": {
"type": "object",
"properties": {
"intent": {
"type": "string",
"description": "Natural-language description of what the user wants \
to build (\"a healthcare flow that handles PHI\", \"a streaming \
chat assistant\", \"a banking endpoint for loan decisions\")."
},
"domain": {
"type": "string",
"enum": [
"generic", "healthcare", "banking", "government", "legal",
"chat", "retrieval", "multi_agent",
"legaltech", "fintech", "pharmatech", "medic_research",
"chat_research", "chat_tools", "chat_skills", "whatsapp",
"voice", "dev", "sales_consultive", "sales_widget",
"workflow_automation", "business_intelligence",
"corporate_integration", "self_learning",
"document_analysis", "ticket_triage",
"content_moderation", "knowledge_extraction",
"compliance_monitoring", "recruitment",
"education", "financial_advisor", "data_pipeline"
],
"description": "Optional explicit domain override. Skips the \
classifier — use when you already know which scaffold you want. \
§Fase 7 ships 33 closed domains: verticals (Fase 7.a), agent \
patterns (Fase 7.b), application patterns (Fase 7.c). See \
`axon://logic/flow_composition` for picking the right one."
}
},
"required": ["intent"],
"additionalProperties": false
}
}),
]
}
pub async fn dispatch_call(
params: Value,
catalog: &Arc<Catalog>,
telemetry: &Arc<Telemetry>,
) -> Result<Value, JsonRpcError> {
let call: ToolCall = serde_json::from_value(params)
.map_err(|e| JsonRpcError::invalid_params(format!("tools/call params: {e}")))?;
let started = Instant::now();
let result = match call.name.as_str() {
"axon.primitives" => primitives(call.arguments, catalog),
"axon.primitive_doc" => primitive_doc(call.arguments, catalog),
"axon.check" => check(call.arguments, telemetry),
"axon.parse" => parse(call.arguments, telemetry),
"axon.compose" => compose_tool(call.arguments, catalog, telemetry),
other => Err(JsonRpcError {
code: -32601,
message: format!("unknown tool: `{other}`"),
data: None,
}),
};
let duration = started.elapsed();
let is_error = match &result {
Err(_) => true,
Ok(v) => v["isError"].as_bool().unwrap_or(false),
};
telemetry.record_tool_call(&call.name, duration, is_error);
result
}
#[derive(Debug, Deserialize)]
struct ToolCall {
name: String,
#[serde(default)]
arguments: Value,
}
#[derive(Debug, Deserialize, Default)]
struct PrimitivesArgs {
#[serde(default)]
category: Option<String>,
}
fn primitives(args: Value, catalog: &Arc<Catalog>) -> Result<Value, JsonRpcError> {
let args: PrimitivesArgs = if args.is_null() {
PrimitivesArgs::default()
} else {
serde_json::from_value(args)
.map_err(|e| JsonRpcError::invalid_params(format!("axon.primitives: {e}")))?
};
let filter = args.category.as_deref().map(parse_category).transpose()?;
let entries: Vec<Value> = catalog
.primitives()
.filter(|p| filter.map(|f| p.category == f).unwrap_or(true))
.map(|p| {
json!({
"name": p.name,
"summary": p.summary,
"category": p.category.as_str(),
"top_level": p.top_level,
"since": p.since,
})
})
.collect();
Ok(mcp_text_result(&serde_json::to_string_pretty(&json!({
"count": entries.len(),
"primitives": entries,
})).unwrap()))
}
fn parse_category(s: &str) -> Result<Category, JsonRpcError> {
match s {
"cognition" => Ok(Category::Cognition),
"cognitive_io" => Ok(Category::CognitiveIo),
"data_plane" => Ok(Category::DataPlane),
"session_types" => Ok(Category::SessionTypes),
"wire" => Ok(Category::Wire),
"operators" => Ok(Category::Operators),
other => Err(JsonRpcError::invalid_params(format!(
"unknown category `{other}` — valid: cognition, cognitive_io, \
data_plane, session_types, wire, operators"
))),
}
}
#[derive(Debug, Deserialize)]
struct PrimitiveDocArgs {
name: String,
}
fn primitive_doc(args: Value, catalog: &Arc<Catalog>) -> Result<Value, JsonRpcError> {
let args: PrimitiveDocArgs = serde_json::from_value(args)
.map_err(|e| JsonRpcError::invalid_params(format!("axon.primitive_doc: {e}")))?;
let prim = catalog.primitive(&args.name).ok_or_else(|| JsonRpcError {
code: -32602,
message: format!(
"unknown primitive `{}` — call axon.primitives to see the catalogue",
args.name
),
data: None,
})?;
let payload = json!({
"name": prim.name,
"summary": prim.summary,
"category": prim.category.as_str(),
"top_level": prim.top_level,
"grammar": prim.grammar,
"since": prim.since,
"body_markdown": prim.body,
});
Ok(json!({
"content": [
{ "type": "text", "text": serde_json::to_string_pretty(&payload).unwrap() }
],
"isError": false,
}))
}
fn mcp_text_result(text: &str) -> Value {
json!({
"content": [
{ "type": "text", "text": text }
],
"isError": false,
})
}
#[derive(Debug, Deserialize)]
struct CheckArgs {
source: String,
#[serde(default)]
filename: Option<String>,
}
fn check(args: Value, telemetry: &Arc<Telemetry>) -> Result<Value, JsonRpcError> {
let args: CheckArgs = serde_json::from_value(args)
.map_err(|e| JsonRpcError::invalid_params(format!("axon.check: {e}")))?;
let filename = args.filename.as_deref().unwrap_or("<axon.check input>");
let outcome = compiler_pipeline::run(&args.source, filename);
let payload = compiler_pipeline::outcome_to_check_payload(&outcome);
let is_error = !payload["ok"].as_bool().unwrap_or(true);
let stage_slug = payload["stage"].as_str().unwrap_or("type_check");
telemetry.record_check(stage_slug, is_error);
Ok(json!({
"content": [
{ "type": "text", "text": serde_json::to_string_pretty(&payload).unwrap() }
],
"isError": is_error,
}))
}
#[derive(Debug, Deserialize)]
struct ParseArgs {
source: String,
#[serde(default)]
filename: Option<String>,
}
fn parse(args: Value, telemetry: &Arc<Telemetry>) -> Result<Value, JsonRpcError> {
let args: ParseArgs = serde_json::from_value(args)
.map_err(|e| JsonRpcError::invalid_params(format!("axon.parse: {e}")))?;
let filename = args.filename.as_deref().unwrap_or("<axon.parse input>");
let outcome = compiler_pipeline::run(&args.source, filename);
let payload = compiler_pipeline::outcome_to_parse_payload(outcome);
let is_error = !payload["ok"].as_bool().unwrap_or(true);
let stage_slug = payload["stage"].as_str().unwrap_or("ir_generate");
telemetry.record_check(stage_slug, is_error);
Ok(json!({
"content": [
{ "type": "text", "text": serde_json::to_string_pretty(&payload).unwrap() }
],
"isError": is_error,
}))
}
#[derive(Debug, Deserialize)]
struct ComposeArgs {
intent: String,
#[serde(default)]
domain: Option<String>,
}
fn compose_tool(
args: Value,
catalog: &Arc<Catalog>,
telemetry: &Arc<Telemetry>,
) -> Result<Value, JsonRpcError> {
let args: ComposeArgs = serde_json::from_value(args)
.map_err(|e| JsonRpcError::invalid_params(format!("axon.compose: {e}")))?;
let domain_override = match args.domain.as_deref() {
Some(s) => match compose::parse_domain_hint(s) {
Some(d) => Some(d),
None => {
return Err(JsonRpcError::invalid_params(format!(
"axon.compose: unknown domain `{s}` — see axon.compose tool \
inputSchema for the closed 33-entry catalog (verticals + \
agent patterns + application patterns + meta-patterns)."
)))
}
},
None => None,
};
let r = compose::compose(&args.intent, domain_override, catalog).map_err(|e| {
JsonRpcError {
code: -32603,
message: format!("axon.compose internal error: {e}"),
data: None,
}
})?;
let is_error = r.axon_check_verdict != "well-formed";
let top_score = r
.alternatives
.first()
.map(|a| a.score)
.unwrap_or(0);
telemetry.record_compose(r.domain.slug(), top_score, domain_override.is_some());
let payload = compose::response_to_json(&r);
Ok(json!({
"content": [
{ "type": "text", "text": serde_json::to_string_pretty(&payload).unwrap() }
],
"isError": is_error,
}))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
fn catalog_with(name: &str, top: bool, cat: Category) -> Arc<Catalog> {
use std::io::Write;
use std::sync::atomic::{AtomicU64, Ordering};
static N: AtomicU64 = AtomicU64::new(0);
let n = N.fetch_add(1, Ordering::Relaxed);
let dir = std::env::temp_dir().join(format!(
"axon-emcp-toolstest-{}-{n}-{name}",
std::process::id(),
));
let _ = std::fs::remove_dir_all(&dir);
let prims = dir.join("primitives");
std::fs::create_dir_all(&prims).unwrap();
let mut f = std::fs::File::create(prims.join(format!("{name}.md"))).unwrap();
let body = format!(
"---\nname: {name}\nsummary: test summary\ncategory: {}\ntop_level: {}\n\
since: Fase X\ngrammar: |\n {name} ...\n---\n\nBody.\n",
cat.as_str(),
top,
);
f.write_all(body.as_bytes()).unwrap();
Arc::new(Catalog::load_from(&dir).unwrap())
}
fn tel() -> Arc<Telemetry> {
Arc::new(Telemetry::new(crate::telemetry::TelemetryConfig {
jsonl_sink: None,
deployment_id: "".into(),
max_samples: 1000,
}))
}
#[tokio::test]
async fn primitives_returns_every_entry_when_unfiltered() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({ "name": "axon.primitives", "arguments": {} }),
&cat, &tel())
.await
.unwrap();
let text = v["content"][0]["text"].as_str().unwrap();
let parsed: Value = serde_json::from_str(text).unwrap();
assert_eq!(parsed["count"], 1);
assert_eq!(parsed["primitives"][0]["name"], "socket");
assert_eq!(parsed["primitives"][0]["top_level"], true);
}
#[tokio::test]
async fn primitives_filters_by_category() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({ "name": "axon.primitives",
"arguments": { "category": "session_types" } }),
&cat, &tel())
.await
.unwrap();
let parsed: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(parsed["count"], 1);
let v2 = dispatch_call(
json!({ "name": "axon.primitives",
"arguments": { "category": "cognition" } }),
&cat, &tel())
.await
.unwrap();
let parsed2: Value =
serde_json::from_str(v2["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(parsed2["count"], 0);
}
#[tokio::test]
async fn primitives_rejects_unknown_category() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let err = dispatch_call(
json!({ "name": "axon.primitives",
"arguments": { "category": "bogus" } }),
&cat, &tel())
.await
.expect_err("must reject");
assert_eq!(err.code, -32602);
assert!(err.message.contains("unknown category"));
}
#[tokio::test]
async fn primitive_doc_returns_full_metadata_and_body() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({ "name": "axon.primitive_doc",
"arguments": { "name": "socket" } }),
&cat, &tel())
.await
.unwrap();
let text = v["content"][0]["text"].as_str().unwrap();
let payload: Value = serde_json::from_str(text).unwrap();
assert_eq!(payload["name"], "socket");
assert_eq!(payload["top_level"], true);
assert!(payload["body_markdown"].as_str().unwrap().contains("Body."));
assert!(payload["grammar"].as_str().unwrap().contains("socket"));
}
#[tokio::test]
async fn primitive_doc_rejects_unknown_name() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let err = dispatch_call(
json!({ "name": "axon.primitive_doc",
"arguments": { "name": "does_not_exist" } }),
&cat, &tel())
.await
.expect_err("must reject");
assert!(err.message.contains("unknown primitive"));
assert!(err.message.contains("axon.primitives"));
}
#[tokio::test]
async fn unknown_tool_returns_method_not_found() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let err = dispatch_call(
json!({ "name": "axon.does_not_exist", "arguments": {} }),
&cat, &tel())
.await
.expect_err("must reject");
assert_eq!(err.code, -32601);
}
#[tokio::test]
async fn check_returns_ok_for_a_well_formed_program() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({
"name": "axon.check",
"arguments": { "source": "persona X { tone: precise }" }
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], false);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["ok"], true);
assert_eq!(payload["errors"].as_array().unwrap().len(), 0);
}
#[tokio::test]
async fn check_returns_diagnostic_with_isError_on_syntax_garbage() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({
"name": "axon.check",
"arguments": { "source": "@@@" }
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], true);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["ok"], false);
let errors = payload["errors"].as_array().unwrap();
assert_eq!(errors.len(), 1);
assert_eq!(errors[0]["severity"], "error");
let stage = errors[0]["stage"].as_str().unwrap();
assert!(matches!(stage, "lex" | "parse"));
assert!(errors[0]["line"].as_u64().unwrap() >= 1);
}
#[tokio::test]
async fn check_rejects_missing_source_argument() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let err = dispatch_call(
json!({ "name": "axon.check", "arguments": {} }),
&cat, &tel())
.await
.expect_err("missing required `source` must reject");
assert_eq!(err.code, -32602);
assert!(err.message.contains("axon.check"));
}
#[tokio::test]
async fn check_accepts_optional_filename_for_diagnostics() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({
"name": "axon.check",
"arguments": {
"source": "persona X { tone: precise }",
"filename": "my_draft.axon"
}
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], false);
}
#[tokio::test]
async fn parse_returns_ir_for_a_well_formed_program() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({
"name": "axon.parse",
"arguments": { "source": "persona X { tone: precise }" }
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], false);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["ok"], true);
assert_eq!(payload["stage"], "ir_generate");
assert_eq!(payload["ir"]["node_type"], "program");
let personas = payload["ir"]["personas"].as_array().unwrap();
assert_eq!(personas.len(), 1);
assert_eq!(personas[0]["name"], "X");
}
#[tokio::test]
async fn parse_returns_same_diagnostic_shape_as_check_on_failure() {
let cat = catalog_with("socket", true, Category::SessionTypes);
let v = dispatch_call(
json!({
"name": "axon.parse",
"arguments": { "source": "@@@" }
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], true);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["ok"], false);
assert!(payload["errors"].as_array().unwrap().len() >= 1);
assert!(payload["ir"].is_null());
}
fn embedded_catalog() -> Arc<Catalog> {
Arc::new(Catalog::load_embedded().expect("embedded corpus must load"))
}
#[tokio::test]
async fn compose_returns_well_formed_scaffold_for_healthcare_intent() {
let cat = embedded_catalog();
let v = dispatch_call(
json!({
"name": "axon.compose",
"arguments": { "intent": "a patient summarisation service with PHI" }
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], false);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["domain"], "healthcare");
assert_eq!(payload["axon_check_verdict"], "well-formed");
assert!(payload["scaffold"].as_str().unwrap().contains("HIPAA"));
let compl: Vec<&str> = payload["compliance_applied"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap())
.collect();
assert!(compl.contains(&"HIPAA"));
}
#[tokio::test]
async fn compose_honors_explicit_domain_argument() {
let cat = embedded_catalog();
let v = dispatch_call(
json!({
"name": "axon.compose",
"arguments": {
"intent": "process credit card payments and loan applications",
"domain": "chat"
}
}),
&cat, &tel())
.await
.unwrap();
assert_eq!(v["isError"], false);
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["domain"], "chat");
assert_eq!(payload["axon_check_verdict"], "well-formed");
}
#[tokio::test]
async fn compose_falls_back_to_generic_for_unrelated_intent() {
let cat = embedded_catalog();
let v = dispatch_call(
json!({
"name": "axon.compose",
"arguments": { "intent": "say hello" }
}),
&cat, &tel())
.await
.unwrap();
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
assert_eq!(payload["domain"], "generic");
assert_eq!(payload["axon_check_verdict"], "well-formed");
}
#[tokio::test]
async fn compose_rejects_unknown_domain_hint() {
let cat = embedded_catalog();
let err = dispatch_call(
json!({
"name": "axon.compose",
"arguments": {
"intent": "anything",
"domain": "bogus-domain"
}
}),
&cat, &tel())
.await
.expect_err("unknown domain must be a structured invalid_params");
assert_eq!(err.code, -32602);
assert!(err.message.contains("unknown domain"));
}
#[tokio::test]
async fn compose_rejects_missing_intent_argument() {
let cat = embedded_catalog();
let err = dispatch_call(
json!({ "name": "axon.compose", "arguments": {} }),
&cat, &tel())
.await
.expect_err("missing required `intent` must reject");
assert_eq!(err.code, -32602);
assert!(err.message.contains("axon.compose"));
}
#[tokio::test]
async fn compose_response_carries_explainability_scoreboard() {
let cat = embedded_catalog();
let v = dispatch_call(
json!({
"name": "axon.compose",
"arguments": { "intent": "patient PHI clinical trial under HIPAA" }
}),
&cat, &tel())
.await
.unwrap();
let payload: Value =
serde_json::from_str(v["content"][0]["text"].as_str().unwrap()).unwrap();
let alts = payload["alternatives"].as_array().unwrap();
assert_eq!(alts.len(), 33);
assert!(alts[0]["score"].as_u64().unwrap() >= 1);
assert_eq!(alts[0]["domain"], "healthcare");
assert!(!payload["next_steps"].as_array().unwrap().is_empty());
assert!(!payload["primitives_used"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn compose_advertised_in_tools_list() {
let names: Vec<String> = list()
.iter()
.map(|v| v["name"].as_str().unwrap().to_string())
.collect();
assert!(
names.contains(&"axon.compose".to_string()),
"axon.compose must be advertised in tools/list; saw {names:?}"
);
}
}