use crate::subagent::control::{Up, send_up};
use crate::subagent::protocol::AgentMsg;
use crate::subagent::replies::{Replies, Reply};
use mcp::inbound::{Answer, Handler, Inbound};
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
pub struct ElicitationBridge {
up: Up,
replies: Arc<Replies>,
cancel: Arc<AtomicBool>,
timeout: Duration,
}
impl ElicitationBridge {
pub fn new(
up: Up,
replies: Arc<Replies>,
cancel: Arc<AtomicBool>,
timeout: Duration,
) -> ElicitationBridge {
ElicitationBridge {
up,
replies,
cancel,
timeout,
}
}
}
impl Handler for ElicitationBridge {
fn handle(&self, req: Inbound) -> Option<Answer> {
let Inbound::Elicit {
message,
requested_schema,
} = req
else {
return None;
};
if self.cancel.load(Ordering::Relaxed) {
return Some(Answer::Cancel);
}
let id = self.replies.next_id();
send_up(
&self.up,
&AgentMsg::ToolRequest {
id,
name: "ask_human".to_string(),
args: json!({
"question": message,
"schema": requested_schema.clone(),
}),
},
);
let deadline = Instant::now() + self.timeout;
match self.replies.wait(id, deadline, &self.cancel) {
Some(Reply::Tool { result, is_error }) => {
if is_error {
return Some(Answer::Cancel);
}
if result
.get("timed_out")
.and_then(Value::as_bool)
.unwrap_or(false)
{
return Some(Answer::Cancel);
}
Some(shape_reply(
result.get("reply").unwrap_or(&Value::Null),
&requested_schema,
))
}
_ => Some(Answer::Cancel),
}
}
}
pub(crate) fn shape_reply(reply: &Value, schema: &Value) -> Answer {
if reply.is_object() {
return Answer::Accept(reply.clone());
}
let text = match reply {
Value::String(s) => s.trim().to_string(),
Value::Null => return Answer::Cancel,
other => other.to_string(),
};
if text.is_empty() {
return Answer::Cancel;
}
if let Ok(Value::Object(m)) = serde_json::from_str::<Value>(&text) {
return Answer::Accept(Value::Object(m));
}
let props = schema.get("properties").and_then(Value::as_object);
if let Some(props) = props
&& props.len() == 1
&& let Some((key, spec)) = props.iter().next()
{
let declared = spec.get("type").and_then(Value::as_str).unwrap_or("string");
if let Some(v) = coerce(&text, declared) {
return Answer::Accept(json!({ key: v }));
}
return Answer::Cancel;
}
Answer::Cancel
}
fn coerce(text: &str, declared: &str) -> Option<Value> {
match declared {
"string" => Some(Value::String(text.to_string())),
"boolean" => match text.to_ascii_lowercase().as_str() {
"true" | "yes" | "y" | "1" | "ok" => Some(Value::Bool(true)),
"false" | "no" | "n" | "0" => Some(Value::Bool(false)),
_ => None,
},
"integer" => text.parse::<i64>().ok().map(|n| json!(n)),
"number" => text
.parse::<f64>()
.ok()
.and_then(|n| serde_json::Number::from_f64(n).map(Value::Number)),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn one_prop(name: &str, ty: &str) -> Value {
json!({"type": "object", "properties": { name: {"type": ty} }})
}
fn accepted(a: Answer) -> Value {
match a {
Answer::Accept(v) => v,
other => panic!("expected Accept, got {other:?}"),
}
}
#[test]
fn a_structured_reply_passes_through() {
let r = shape_reply(&json!({"env": "staging"}), &one_prop("env", "string"));
assert_eq!(accepted(r)["env"], "staging");
}
#[test]
fn a_single_property_schema_binds_the_text() {
let r = shape_reply(&json!("staging"), &one_prop("env", "string"));
assert_eq!(accepted(r), json!({"env": "staging"}));
let r = shape_reply(&json!("42"), &one_prop("count", "integer"));
assert_eq!(accepted(r), json!({"count": 42}));
let r = shape_reply(&json!("yes"), &one_prop("confirm", "boolean"));
assert_eq!(accepted(r), json!({"confirm": true}));
assert!(matches!(
shape_reply(&json!("soon"), &one_prop("count", "integer")),
Answer::Cancel
));
}
#[test]
fn typed_json_is_honoured_over_the_single_property_shortcut() {
let r = shape_reply(
&json!(r#"{"env":"prod","force":true}"#),
&one_prop("env", "string"),
);
let v = accepted(r);
assert_eq!(v["env"], "prod");
assert_eq!(v["force"], true);
}
#[test]
fn prose_against_a_multi_property_schema_is_declined_not_guessed() {
let schema = json!({"type": "object", "properties": {
"env": {"type": "string"}, "force": {"type": "boolean"}
}});
assert!(matches!(
shape_reply(&json!("just do it on staging"), &schema),
Answer::Cancel
));
}
#[test]
fn nothing_to_say_is_a_cancel() {
assert!(matches!(
shape_reply(&Value::Null, &one_prop("x", "string")),
Answer::Cancel
));
assert!(matches!(
shape_reply(&json!(" "), &one_prop("x", "string")),
Answer::Cancel
));
}
}