use crate::{BehaviorNode, Context, NodeStatus, TreeStatusSnapshot};
use async_trait::async_trait;
use mecha10_core::Receiver;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TargetRequest {
#[serde(default)]
bbox: Option<Value>,
}
#[derive(Debug, Serialize, Deserialize)]
struct CapabilityTaskConfig {
capability: String,
task_complete_topic: String,
#[serde(default)]
target_request_topic: Option<String>,
}
pub struct CapabilityTaskNode {
capability: String,
task_complete_topic: String,
target_request_topic: Option<String>,
complete_receiver: Option<Receiver<Value>>,
target_receiver: Option<Receiver<Value>>,
complete: bool,
has_target: bool,
last_status: Option<NodeStatus>,
}
impl std::fmt::Debug for CapabilityTaskNode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CapabilityTaskNode")
.field("capability", &self.capability)
.field("task_complete_topic", &self.task_complete_topic)
.field("target_request_topic", &self.target_request_topic)
.field("complete_subscribed", &self.complete_receiver.is_some())
.field("target_subscribed", &self.target_receiver.is_some())
.field("complete", &self.complete)
.field("has_target", &self.has_target)
.finish()
}
}
impl CapabilityTaskNode {
pub fn from_json(config: Value) -> anyhow::Result<Self> {
let cfg: CapabilityTaskConfig = serde_json::from_value(config)?;
Ok(Self {
capability: cfg.capability,
task_complete_topic: cfg.task_complete_topic,
target_request_topic: cfg.target_request_topic,
complete_receiver: None,
target_receiver: None,
complete: false,
has_target: false,
last_status: None,
})
}
#[cfg(test)]
fn with_receivers(
capability: &str,
task_complete_topic: &str,
target_request_topic: Option<&str>,
complete_receiver: Receiver<Value>,
target_receiver: Option<Receiver<Value>>,
) -> Self {
Self {
capability: capability.to_string(),
task_complete_topic: task_complete_topic.to_string(),
target_request_topic: target_request_topic.map(|s| s.to_string()),
complete_receiver: Some(complete_receiver),
target_receiver,
complete: false,
has_target: false,
last_status: None,
}
}
fn poll_target(&mut self) {
let Some(receiver) = self.target_receiver.as_mut() else {
return;
};
let mut latest = None;
while let Ok(value) = receiver.try_recv() {
match serde_json::from_value::<TargetRequest>(value) {
Ok(request) => latest = Some(request.bbox.is_some()),
Err(e) => tracing::warn!("CapabilityTaskNode: dropping malformed target request: {}", e),
}
}
if let Some(has_target) = latest {
self.has_target = has_target;
}
}
fn poll_complete(&mut self) {
if self.target_request_topic.is_none() && self.complete {
return;
}
let Some(receiver) = self.complete_receiver.as_mut() else {
return;
};
if receiver.try_recv().is_ok() {
self.complete = true;
self.has_target = false;
}
}
}
#[async_trait]
impl BehaviorNode for CapabilityTaskNode {
async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
if self.complete_receiver.is_none() {
self.complete_receiver = Some(ctx.subscribe_raw::<Value>(&self.task_complete_topic).await?);
}
if let Some(target_request_topic) = self.target_request_topic.clone() {
if self.target_receiver.is_none() {
self.target_receiver = Some(ctx.subscribe_raw::<Value>(&target_request_topic).await?);
}
}
let status = if self.target_request_topic.is_some() {
self.poll_target();
self.poll_complete();
if self.has_target {
NodeStatus::Running
} else {
NodeStatus::Failure
}
} else {
self.poll_complete();
if self.complete {
NodeStatus::Success
} else {
NodeStatus::Running
}
};
self.last_status = Some(status);
Ok(status)
}
fn name(&self) -> &str {
"capability_task"
}
fn active_capability(&self) -> Option<&str> {
Some(self.capability.as_str())
}
fn reset_self(&mut self) {
self.complete = false;
self.has_target = false;
}
fn status_snapshot(&self) -> TreeStatusSnapshot {
TreeStatusSnapshot {
name: self.name().to_string(),
status: self.last_status,
active_capability: self.active_capability().map(String::from),
children: Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::mpsc;
fn task_only_node() -> (CapabilityTaskNode, mpsc::UnboundedSender<Value>) {
let (tx, rx) = mpsc::unbounded_channel();
(
CapabilityTaskNode::with_receivers("rl", "/rl/task_complete", None, Receiver::from(rx), None),
tx,
)
}
fn tracking_node() -> (
CapabilityTaskNode,
mpsc::UnboundedSender<Value>,
mpsc::UnboundedSender<Value>,
) {
let (complete_tx, complete_rx) = mpsc::unbounded_channel();
let (target_tx, target_rx) = mpsc::unbounded_channel();
(
CapabilityTaskNode::with_receivers(
"instance-track",
"/instance-track/task_complete",
Some("/instance-track/target/request"),
Receiver::from(complete_rx),
Some(Receiver::from(target_rx)),
),
complete_tx,
target_tx,
)
}
#[test]
fn declares_the_configured_capability() {
let (node, _tx) = task_only_node();
assert_eq!(node.active_capability(), Some("rl"));
let (node, _complete_tx, _target_tx) = tracking_node();
assert_eq!(node.active_capability(), Some("instance-track"));
}
#[test]
fn task_only_stays_incomplete_with_no_message() {
let (mut node, _tx) = task_only_node();
node.poll_complete();
assert!(!node.complete);
}
#[test]
fn task_only_completes_once_any_message_arrives() {
let (mut node, tx) = task_only_node();
tx.send(serde_json::json!({})).unwrap();
node.poll_complete();
assert!(node.complete);
}
#[test]
fn task_only_reset_self_clears_completion_for_a_later_pass() {
let (mut node, tx) = task_only_node();
tx.send(serde_json::json!({})).unwrap();
node.poll_complete();
assert!(node.complete);
node.reset_self();
assert!(!node.complete);
}
#[test]
fn tracking_no_target_ever_set_stays_without_a_target() {
let (mut node, _complete_tx, _target_tx) = tracking_node();
node.poll_target();
node.poll_complete();
assert!(!node.has_target);
}
#[test]
fn tracking_a_live_target_request_sets_has_target() {
let (mut node, _complete_tx, target_tx) = tracking_node();
target_tx
.send(serde_json::json!({ "bbox": { "x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4 } }))
.unwrap();
node.poll_target();
assert!(node.has_target);
}
#[test]
fn tracking_a_clear_request_unsets_has_target_even_with_no_task_complete() {
let (mut node, _complete_tx, target_tx) = tracking_node();
target_tx
.send(serde_json::json!({ "bbox": { "x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4 } }))
.unwrap();
node.poll_target();
assert!(node.has_target);
target_tx.send(serde_json::json!({ "bbox": null })).unwrap();
node.poll_target();
assert!(!node.has_target);
}
#[test]
fn tracking_task_complete_unsets_has_target_on_the_very_next_poll() {
let (mut node, complete_tx, target_tx) = tracking_node();
target_tx
.send(serde_json::json!({ "bbox": { "x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4 } }))
.unwrap();
node.poll_target();
assert!(node.has_target);
complete_tx.send(serde_json::json!({})).unwrap();
node.poll_complete();
assert!(
!node.has_target,
"task_complete must converge has_target to false without waiting for a separate clear request"
);
}
#[test]
fn tracking_reset_self_clears_has_target_for_a_later_pass() {
let (mut node, _complete_tx, target_tx) = tracking_node();
target_tx
.send(serde_json::json!({ "bbox": { "x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4 } }))
.unwrap();
node.poll_target();
assert!(node.has_target);
node.reset_self();
assert!(!node.has_target);
}
#[test]
fn from_json_requires_capability_and_task_complete_topic() {
let node = CapabilityTaskNode::from_json(serde_json::json!({
"capability": "rl",
"task_complete_topic": "/rl/task_complete"
}))
.unwrap();
assert_eq!(node.capability, "rl");
assert_eq!(node.task_complete_topic, "/rl/task_complete");
assert_eq!(node.target_request_topic, None);
assert!(CapabilityTaskNode::from_json(serde_json::json!({ "capability": "rl" })).is_err());
assert!(CapabilityTaskNode::from_json(serde_json::json!({})).is_err());
}
#[test]
fn from_json_reads_optional_target_request_topic() {
let node = CapabilityTaskNode::from_json(serde_json::json!({
"capability": "instance-track",
"task_complete_topic": "/instance-track/task_complete",
"target_request_topic": "/instance-track/target/request"
}))
.unwrap();
assert_eq!(node.capability, "instance-track");
assert_eq!(
node.target_request_topic,
Some("/instance-track/target/request".to_string())
);
}
}