Add procedure-scoped stateful leaves

src/protocol/tests/mod.rs

@@ -1,3 +1,4 @@
 mod call;
+mod procedure;
 mod protocol;
 mod tree;

src/protocol/tests/procedure.rs

@@ -0,0 +1,267 @@
+use alloc::{borrow::ToOwned, collections::BTreeMap, format, string::String, vec, vec::Vec};
+use core::convert::Infallible;
+
+use crate::protocol::tree::{
+    Call, ChildRoute, ConnectionState, Endpoint, HookKey, Ingress, OutgoingData, Procedure,
+    ProcedureEffect, ProcedureRuntime, ProcedureStore, ProtocolEndpoint, encode_call_reply,
+};
+use crate::protocol::{PacketType, decode_frame};
+use crate::{Leaf, Procedure};
+
+fn path(parts: &[&str]) -> Vec<String> {
+    parts.iter().map(|part| (*part).to_owned()).collect()
+}
+
+#[derive(Default, Leaf)]
+#[leaf(id = "org.example.v1.stream")]
+struct StreamLeaf {
+    sessions: BTreeMap<HookKey, ProcedureOpen>,
+}
+
+impl ProcedureStore<ProcedureOpen> for StreamLeaf {
+    fn procedure_sessions(&mut self) -> &mut BTreeMap<HookKey, ProcedureOpen> {
+        &mut self.sessions
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Procedure)]
+#[procedure(leaf = StreamLeaf, name = "open")]
+struct ProcedureOpen {
+    prefix: String,
+}
+
+impl Procedure<StreamLeaf> for ProcedureOpen {
+    type Error = Infallible;
+    type Input = String;
+
+    fn open(_leaf: &mut StreamLeaf, call: Call<Self::Input>) -> Result<Self, Self::Error> {
+        Ok(Self { prefix: call.input })
+    }
+
+    fn on_data(
+        _leaf: &mut StreamLeaf,
+        session: &mut Self,
+        data: crate::protocol::tree::IncomingData,
+    ) -> Result<ProcedureEffect, Self::Error> {
+        Ok(ProcedureEffect {
+            outgoing: vec![OutgoingData {
+                dst_path: data.hook_key.return_path,
+                hook_id: data.hook_key.hook_id,
+                procedure_id: ProcedureOpen::protocol_procedure_id(),
+                data: format!(
+                    "{}{}",
+                    session.prefix,
+                    String::from_utf8_lossy(&data.message.data)
+                )
+                .into_bytes(),
+                end_hook: data.message.end_hook,
+            }],
+            close_session: data.message.end_hook,
+        })
+    }
+}
+
+#[test]
+fn procedure_runtime_routes_data_to_stored_session() {
+    let endpoint = ProtocolEndpoint::new(
+        path(&["agent"]),
+        Some(Vec::new()),
+        Vec::new(),
+        vec![crate::protocol::tree::LeafSpec {
+            name: StreamLeaf::protocol_leaf_name(),
+            procedures: vec![ProcedureOpen::protocol_procedure_id()],
+        }],
+    );
+    let mut runtime =
+        ProcedureRuntime::<StreamLeaf, ProcedureOpen>::new(endpoint, StreamLeaf::default());
+
+    let mut controller = ProtocolEndpoint::new(
+        Vec::new(),
+        None,
+        vec![ChildRoute {
+            path: path(&["agent"]),
+            state: ConnectionState::Registered,
+        }],
+        Vec::new(),
+    );
+    let hook_id = controller.allocate_hook_id();
+    let open = controller
+        .send_call(
+            path(&["agent"]),
+            Some(StreamLeaf::protocol_leaf_name()),
+            ProcedureOpen::protocol_procedure_id(),
+            Some(hook_id),
+            encode_call_reply(&String::from("prefix:")).expect("procedure input should encode"),
+        )
+        .expect("open call should encode");
+    let Some((_, open_frame)) = open.forward else {
+        panic!("controller should forward opening call");
+    };
+    runtime
+        .receive(&Ingress::Parent, open_frame)
+        .expect("runtime should open a session");
+
+    let data = controller
+        .send_data(
+            path(&["agent"]),
+            hook_id,
+            ProcedureOpen::protocol_procedure_id(),
+            b"hello".to_vec(),
+            true,
+        )
+        .expect("data should encode");
+    let Some((_, data_frame)) = data.forward else {
+        panic!("controller should forward data frame");
+    };
+    let outcome = runtime
+        .receive(&Ingress::Parent, data_frame)
+        .expect("runtime should route data to session");
+    let [response_frame] = outcome.frames.as_slice() else {
+        panic!("expected one response frame");
+    };
+
+    let parsed = decode_frame(response_frame.as_slice()).expect("response frame should decode");
+    assert_eq!(parsed.packet_type(), PacketType::Data);
+    let message = parsed.deserialize_data().expect("data should deserialize");
+    assert!(message.end_hook);
+    assert_eq!(String::from_utf8_lossy(&message.data), "prefix:hello");
+
+    let forwarded = controller
+        .receive(&Ingress::Child(path(&["agent"])), response_frame.clone())
+        .expect("controller should receive session response");
+    assert!(forwarded.event.is_some());
+    assert!(runtime.leaf_mut().procedure_sessions().is_empty());
+}
+
+#[derive(Default, Leaf)]
+#[leaf(id = "org.example.v1.duplex")]
+struct DuplexLeaf {
+    sessions: BTreeMap<HookKey, DuplexProcedure>,
+}
+
+impl ProcedureStore<DuplexProcedure> for DuplexLeaf {
+    fn procedure_sessions(&mut self) -> &mut BTreeMap<HookKey, DuplexProcedure> {
+        &mut self.sessions
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Procedure)]
+#[procedure(leaf = DuplexLeaf, name = "open")]
+struct DuplexProcedure {
+    saw_peer_close: bool,
+}
+
+impl Procedure<DuplexLeaf> for DuplexProcedure {
+    type Error = Infallible;
+    type Input = ();
+
+    fn open(_leaf: &mut DuplexLeaf, _call: Call<Self::Input>) -> Result<Self, Self::Error> {
+        Ok(Self {
+            saw_peer_close: false,
+        })
+    }
+
+    fn on_data(
+        _leaf: &mut DuplexLeaf,
+        session: &mut Self,
+        data: crate::protocol::tree::IncomingData,
+    ) -> Result<ProcedureEffect, Self::Error> {
+        if data.message.data == b"local-end" {
+            return Ok(ProcedureEffect::outgoing(vec![OutgoingData {
+                dst_path: data.hook_key.return_path,
+                hook_id: data.hook_key.hook_id,
+                procedure_id: DuplexProcedure::protocol_procedure_id(),
+                data: Vec::new(),
+                end_hook: true,
+            }]));
+        }
+
+        if data.message.end_hook {
+            session.saw_peer_close = true;
+            return Ok(ProcedureEffect::close(Vec::new()));
+        }
+
+        Ok(ProcedureEffect::default())
+    }
+}
+
+#[test]
+fn procedure_runtime_keeps_session_after_local_end_until_explicit_close() {
+    let endpoint = ProtocolEndpoint::new(
+        path(&["agent"]),
+        Some(Vec::new()),
+        Vec::new(),
+        vec![crate::protocol::tree::LeafSpec {
+            name: DuplexLeaf::protocol_leaf_name(),
+            procedures: vec![DuplexProcedure::protocol_procedure_id()],
+        }],
+    );
+    let mut runtime =
+        ProcedureRuntime::<DuplexLeaf, DuplexProcedure>::new(endpoint, DuplexLeaf::default());
+
+    let mut controller = ProtocolEndpoint::new(
+        Vec::new(),
+        None,
+        vec![ChildRoute {
+            path: path(&["agent"]),
+            state: ConnectionState::Registered,
+        }],
+        Vec::new(),
+    );
+    let hook_id = controller.allocate_hook_id();
+    let open = controller
+        .send_call(
+            path(&["agent"]),
+            Some(DuplexLeaf::protocol_leaf_name()),
+            DuplexProcedure::protocol_procedure_id(),
+            Some(hook_id),
+            encode_call_reply(&()).expect("unit call should encode"),
+        )
+        .expect("open call should encode");
+    let Some((_, open_frame)) = open.forward else {
+        panic!("controller should forward opening call");
+    };
+    runtime
+        .receive(&Ingress::Parent, open_frame)
+        .expect("runtime should open duplex session");
+
+    let local_end = controller
+        .send_data(
+            path(&["agent"]),
+            hook_id,
+            DuplexProcedure::protocol_procedure_id(),
+            b"local-end".to_vec(),
+            false,
+        )
+        .expect("local end trigger should encode");
+    let Some((_, local_end_frame)) = local_end.forward else {
+        panic!("controller should forward local end trigger");
+    };
+    let outcome = runtime
+        .receive(&Ingress::Parent, local_end_frame)
+        .expect("runtime should emit a local end packet");
+    assert_eq!(outcome.frames.len(), 1);
+    assert_eq!(runtime.leaf_mut().procedure_sessions().len(), 1);
+
+    let peer_end = encode_call_reply(&()).expect("unit value is just a placeholder");
+    let peer_end = crate::protocol::encode_packet(
+        &crate::protocol::PacketHeader {
+            packet_type: PacketType::Data,
+            src_path: Vec::new(),
+            dst_path: path(&["agent"]),
+            dst_leaf: None,
+            hook_id: Some(hook_id),
+        },
+        &crate::protocol::DataMessage {
+            procedure_id: DuplexProcedure::protocol_procedure_id(),
+            data: peer_end,
+            end_hook: true,
+        },
+    )
+    .expect("peer end frame should encode");
+    let peer_end_outcome = runtime
+        .receive(&Ingress::Parent, peer_end)
+        .expect("runtime should accept peer end after local end");
+    assert!(peer_end_outcome.frames.is_empty());
+    assert!(runtime.leaf_mut().procedure_sessions().is_empty());
+}

src/protocol/tests/tree.rs

@@ -106,7 +106,15 @@ fn protocol_endpoint_introspection_returns_leaf_summary() {
 
 #[test]
 fn invalid_hook_peer_emits_local_fault_event() {
-    let mut endpoint = ProtocolEndpoint::new(path(&["client"]), None, Vec::new(), Vec::new());
+    let mut endpoint = ProtocolEndpoint::new(
+        Vec::new(),
+        None,
+        vec![
+            ChildRoute::registered(path(&["server"])),
+            ChildRoute::registered(path(&["intruder"])),
+        ],
+        Vec::new(),
+    );
     let hook_id = endpoint.allocate_hook_id();
 
     endpoint
@@ -119,11 +127,31 @@ fn invalid_hook_peer_emits_local_fault_event() {
         )
         .expect("call should establish an active hook");
 
+    let valid_frame = encode_packet(
+        &PacketHeader {
+            packet_type: PacketType::Data,
+            src_path: path(&["server"]),
+            dst_path: Vec::new(),
+            dst_leaf: None,
+            hook_id: Some(hook_id),
+        },
+        &DataMessage {
+            procedure_id: "example.service.v1.invoke".to_owned(),
+            data: vec![8],
+            end_hook: false,
+        },
+    )
+    .expect("valid server data should encode");
+
+    endpoint
+        .receive(&Ingress::Child(path(&["server"])), valid_frame)
+        .expect("first server data should activate the hook");
+
     let frame = encode_packet(
         &PacketHeader {
             packet_type: PacketType::Data,
-            src_path: path(&["client"]),
-            dst_path: path(&["client"]),
+            src_path: path(&["intruder"]),
+            dst_path: Vec::new(),
             dst_leaf: None,
             hook_id: Some(hook_id),
         },
@@ -136,13 +164,13 @@ fn invalid_hook_peer_emits_local_fault_event() {
     .expect("data frame should encode");
 
     let outcome = endpoint
-        .receive(&Ingress::Local, frame)
+        .receive(&Ingress::Child(path(&["intruder"])), frame)
         .expect("invalid peer should be handled");
 
     assert!(outcome.forward.is_none());
     assert!(!outcome.dropped);
 
-    match outcome.event.as_ref().expect("expected event") {
+    match outcome.event.as_ref().expect("expected local fault event") {
         LocalEvent::Fault {
             header, message, ..
         } => {
@@ -180,6 +208,27 @@ fn hook_closes_only_after_both_sides_end() {
         .expect("call should establish an active hook");
 
     let host_key = crate::protocol::tree::HookKey::new(Vec::new(), hook_id);
+    assert!(endpoint.hooks.pending(&host_key).is_some());
+
+    let activation_frame = encode_packet(
+        &PacketHeader {
+            packet_type: PacketType::Data,
+            src_path: path(&["server"]),
+            dst_path: Vec::new(),
+            dst_leaf: None,
+            hook_id: Some(hook_id),
+        },
+        &DataMessage {
+            procedure_id: "example.service.v1.invoke".to_owned(),
+            data: vec![9],
+            end_hook: false,
+        },
+    )
+    .expect("activation data should encode");
+
+    endpoint
+        .receive(&Ingress::Child(path(&["server"])), activation_frame)
+        .expect("first server data should activate the hook");
     assert!(endpoint.hooks.active(&host_key).is_some());
 
     endpoint
@@ -242,6 +291,7 @@ fn pending_hook_fault_is_delivered_before_activation() {
             caller_src_path: path(&["client"]),
             procedure_id: call.procedure_id.clone(),
             dst_leaf: None,
+            local_ended: false,
         })
         .expect("pending hook should insert");