add realistic root knowledge mode to treetest

This commit is contained in:
Michael Mikovsky
2026-04-24 16:46:56 -06:00
parent de3a7d3381
commit 18cdefaefb
5 changed files with 656 additions and 76 deletions
+42 -9
View File
@@ -11,7 +11,11 @@ use crossterm::{
}; };
use ratatui::DefaultTerminal; use ratatui::DefaultTerminal;
use crate::{model::Selection, scenarios::built_in_scenarios, sim::Simulation}; use crate::{
model::{NodeId, Selection},
scenarios::built_in_scenarios,
sim::Simulation,
};
/// Errors returned by the TUI application. /// Errors returned by the TUI application.
#[derive(Debug, thiserror::Error)] #[derive(Debug, thiserror::Error)]
@@ -113,6 +117,21 @@ impl App {
KeyCode::Char('d') => self.perform_chat_data()?, KeyCode::Char('d') => self.perform_chat_data()?,
KeyCode::Char('b') => self.perform_chat_bye()?, KeyCode::Char('b') => self.perform_chat_bye()?,
KeyCode::Char('f') => self.perform_invalid_fault_demo()?, KeyCode::Char('f') => self.perform_invalid_fault_demo()?,
KeyCode::Char('g') => {
self.simulation.toggle_inspector_mode();
self.refresh_selections(Some(self.selected().node_id()));
self.status = if self.simulation.is_realistic_mode() {
"Inspector switched to realistic mode.".to_owned()
} else {
"Inspector switched to ground truth mode.".to_owned()
};
}
KeyCode::Char('m') => {
self.simulation.enable_realistic_mode_with_memory_reset();
self.refresh_selections(Some(NodeId(0)));
self.status =
"Cleared root memory for deeper nodes and enabled realistic mode.".to_owned();
}
KeyCode::Char('s') => { KeyCode::Char('s') => {
let processed = self.simulation.step()?; let processed = self.simulation.step()?;
self.status = if processed { self.status = if processed {
@@ -133,12 +152,7 @@ impl App {
fn load_scenario(&mut self, index: usize) -> Result<(), AppError> { fn load_scenario(&mut self, index: usize) -> Result<(), AppError> {
self.scenario_index = index; self.scenario_index = index;
self.simulation = Simulation::new(self.scenarios[index].clone())?; self.simulation = Simulation::new(self.scenarios[index].clone())?;
self.selections = ui::build_selections(&self.simulation); self.refresh_selections(Some(self.simulation.initial_selection().node_id()));
self.selection_index = self
.selections
.iter()
.position(|selection| *selection == self.simulation.initial_selection())
.unwrap_or(0);
self.status = format!("Loaded scenario: {}", self.scenarios[index].name); self.status = format!("Loaded scenario: {}", self.scenarios[index].name);
Ok(()) Ok(())
} }
@@ -152,6 +166,7 @@ impl App {
Selection::Node(node_id) => { Selection::Node(node_id) => {
let result = self.simulation.call_endpoint_introspection(node_id)?; let result = self.simulation.call_endpoint_introspection(node_id)?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} }
Selection::Leaf { node_id, leaf_name } => { Selection::Leaf { node_id, leaf_name } => {
@@ -159,6 +174,7 @@ impl App {
.simulation .simulation
.call_leaf_introspection(node_id, &leaf_name)?; .call_leaf_introspection(node_id, &leaf_name)?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} }
} }
@@ -171,6 +187,7 @@ impl App {
self.simulation self.simulation
.call_echo_leaf(node_id, &leaf_name, "demo echo from root")?; .call_echo_leaf(node_id, &leaf_name, "demo echo from root")?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "Select a leaf first, then press e.".to_owned(); self.status = "Select a leaf first, then press e.".to_owned();
@@ -193,6 +210,7 @@ impl App {
b"ping".to_vec(), b"ping".to_vec(),
)?; )?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "Selected node has no endpoint procedures.".to_owned(); self.status = "Selected node has no endpoint procedures.".to_owned();
@@ -222,6 +240,7 @@ impl App {
b"chunk please".to_vec(), b"chunk please".to_vec(),
)?; )?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "Selected node has no chunked procedure.".to_owned(); self.status = "Selected node has no chunked procedure.".to_owned();
@@ -248,6 +267,7 @@ impl App {
b"open chat".to_vec(), b"open chat".to_vec(),
)?; )?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(Some(node_id));
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "Selected node has no chat procedure.".to_owned(); self.status = "Selected node has no chat procedure.".to_owned();
@@ -264,6 +284,7 @@ impl App {
self.simulation self.simulation
.send_root_hook_data(hook_id, "hello from the root", false)?; .send_root_hook_data(hook_id, "hello from the root", false)?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(None);
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "No known hook yet. Press h to open chat first.".to_owned(); self.status = "No known hook yet. Press h to open chat first.".to_owned();
@@ -275,6 +296,7 @@ impl App {
if let Some(hook_id) = self.simulation.hook_ids().last().copied() { if let Some(hook_id) = self.simulation.hook_ids().last().copied() {
let result = self.simulation.send_root_hook_data(hook_id, "bye", true)?; let result = self.simulation.send_root_hook_data(hook_id, "bye", true)?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(None);
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = "No known hook yet. Press h to open chat first.".to_owned(); self.status = "No known hook yet. Press h to open chat first.".to_owned();
@@ -284,9 +306,9 @@ impl App {
fn perform_invalid_fault_demo(&mut self) -> Result<(), AppError> { fn perform_invalid_fault_demo(&mut self) -> Result<(), AppError> {
if let Some(hook_id) = self.simulation.hook_ids().last().copied() { if let Some(hook_id) = self.simulation.hook_ids().last().copied() {
let root_id = crate::model::NodeId(0); let root_id = NodeId(0);
if self.simulation.tree.nodes.len() > 1 { if self.simulation.tree.nodes.len() > 1 {
let attacker = crate::model::NodeId(1); let attacker = NodeId(1);
let result = self.simulation.inject_invalid_peer_data( let result = self.simulation.inject_invalid_peer_data(
attacker, attacker,
root_id, root_id,
@@ -295,6 +317,7 @@ impl App {
"spoofed data", "spoofed data",
)?; )?;
let steps = self.simulation.drain()?; let steps = self.simulation.drain()?;
self.refresh_selections(None);
self.status = format!("{} ({steps} steps)", result.label); self.status = format!("{} ({steps} steps)", result.label);
} else { } else {
self.status = self.status =
@@ -305,4 +328,14 @@ impl App {
} }
Ok(()) Ok(())
} }
fn refresh_selections(&mut self, preferred_node: Option<NodeId>) {
let current = preferred_node.unwrap_or_else(|| self.selected().node_id());
self.selections = ui::build_selections(&self.simulation);
self.selection_index = self
.selections
.iter()
.position(|selection| selection.node_id() == current)
.unwrap_or(0);
}
} }
+174 -38
View File
@@ -9,8 +9,8 @@ use ratatui::{
}; };
use crate::{ use crate::{
model::{Selection, format_path}, model::{Selection, format_hook_ref, format_leaf_ref, format_path},
sim::{RecordedEvent, Simulation}, sim::{InspectorMode, RecordedEvent, Simulation},
}; };
use super::App; use super::App;
@@ -32,11 +32,16 @@ impl App {
} }
fn render_header(&self, frame: &mut Frame<'_>, area: Rect) { fn render_header(&self, frame: &mut Frame<'_>, area: Rect) {
let mode = match self.simulation.inspector_mode {
InspectorMode::GroundTruth => "ground truth",
InspectorMode::Realistic => "realistic",
};
let title = format!( let title = format!(
"treetest | scenario {} / {}: {}", "treetest | scenario {} / {}: {} | {}",
self.scenario_index + 1, self.scenario_index + 1,
self.scenarios.len(), self.scenarios.len(),
self.scenarios[self.scenario_index].name self.scenarios[self.scenario_index].name,
mode
); );
frame.render_widget( frame.render_widget(
Paragraph::new(title).block(Block::default().borders(Borders::ALL).title("Scenario")), Paragraph::new(title).block(Block::default().borders(Borders::ALL).title("Scenario")),
@@ -107,18 +112,15 @@ impl App {
node.display_path() node.display_path()
) )
} }
Selection::Leaf { node_id, leaf_name } => { Selection::Leaf { node_id, leaf_name } => format!(
format!( "{} {}",
"{} {} :: {}",
if index == self.selection_index { if index == self.selection_index {
">" ">"
} else { } else {
" " " "
}, },
self.simulation.node(*node_id).display_path(), format_leaf_ref(&self.simulation.node(*node_id).path, leaf_name)
leaf_name ),
)
}
}; };
ListItem::new(label) ListItem::new(label)
}) })
@@ -131,7 +133,21 @@ impl App {
fn render_inspector(&self, frame: &mut Frame<'_>, area: Rect) { fn render_inspector(&self, frame: &mut Frame<'_>, area: Rect) {
let selection = self.selected(); let selection = self.selected();
let body = match selection { let body = match self.simulation.inspector_mode {
InspectorMode::GroundTruth => self.render_ground_truth_inspector(selection),
InspectorMode::Realistic => self.render_realistic_inspector(selection),
};
frame.render_widget(
Paragraph::new(body)
.block(Block::default().borders(Borders::ALL).title("Inspector"))
.wrap(Wrap { trim: true }),
area,
);
}
fn render_ground_truth_inspector(&self, selection: &Selection) -> Text<'static> {
match selection {
Selection::Node(node_id) => { Selection::Node(node_id) => {
let node = self.simulation.node(*node_id); let node = self.simulation.node(*node_id);
let mut lines = vec![ let mut lines = vec![
@@ -147,16 +163,15 @@ impl App {
Line::default(), Line::default(),
Line::from("Endpoint procedures:"), Line::from("Endpoint procedures:"),
]; ];
lines.extend( lines.extend(node.endpoint_procedures.iter().map(|procedure| {
node.endpoint_procedures Line::from(format!(
.iter() "- {}: {}",
.map(|procedure| Line::from(format!("- {}", procedure.procedure_id))), procedure.procedure_id, procedure.description
); ))
lines.extend( }));
node.leaves lines.extend(node.leaves.iter().map(|leaf| {
.iter() Line::from(format!("- {}", format_leaf_ref(&node.path, &leaf.name)))
.map(|leaf| Line::from(format!("- leaf {}", leaf.name))), }));
);
Text::from(lines) Text::from(lines)
} }
Selection::Leaf { node_id, leaf_name } => { Selection::Leaf { node_id, leaf_name } => {
@@ -166,21 +181,118 @@ impl App {
.iter() .iter()
.find(|leaf| &leaf.name == leaf_name) .find(|leaf| &leaf.name == leaf_name)
.expect("selection should stay valid"); .expect("selection should stay valid");
Text::from(vec![ let mut lines = vec![
Line::from(format!("Leaf {}", leaf.name)).bold(), Line::from(format_leaf_ref(&node.path, &leaf.name)).bold(),
Line::from(leaf.description.clone()), Line::from(leaf.description.clone()),
Line::from(format!("Node: {}", node.display_path())), Line::from(format!("Node: {}", node.display_path())),
Line::from(format!("Procedures: {}", leaf.procedures.join(", "))), Line::from("Procedures:"),
];
lines.extend(
leaf.procedures
.iter()
.map(|procedure| Line::from(format!("- {}", procedure))),
);
Text::from(lines)
}
}
}
fn render_realistic_inspector(&self, selection: &Selection) -> Text<'static> {
match selection {
Selection::Node(node_id) => {
let node = self.simulation.node(*node_id);
if let Some(learned) = self.simulation.root_knowledge.node(&node.path) {
let mut lines = vec![
Line::from(learned.title.clone().unwrap_or_else(|| node.display_path()))
.bold(),
Line::from(
learned
.description
.clone()
.unwrap_or_else(|| "No learned description yet.".to_owned()),
),
Line::from(format!("Path: {}", format_path(&learned.path))),
Line::from(format!("Known direct child: {}", learned.direct_child)),
Line::from(format!(
"Endpoint introspected: {}",
learned.endpoint_introspected
)),
Line::default(),
Line::from("Known endpoint procedures:"),
];
if learned.endpoint_procedures.is_empty() {
lines.push(Line::from("- none learned"));
} else {
lines.extend(learned.endpoint_procedures.iter().map(|procedure| {
Line::from(match &procedure.description {
Some(description) => {
format!("- {}: {}", procedure.procedure_id, description)
}
None => format!("- {}", procedure.procedure_id),
})
}));
}
lines.push(Line::default());
lines.push(Line::from("Known leaves:"));
if learned.leaves.is_empty() {
lines.push(Line::from("- none learned"));
} else {
lines.extend(learned.leaves.iter().map(|leaf| {
Line::from(format!(
"- {}",
format_leaf_ref(&learned.path, &leaf.leaf_name)
))
}));
}
Text::from(lines)
} else {
Text::from(vec![
Line::from(node.display_path()).bold(),
Line::from(
"The root host has not learned anything about this endpoint yet.",
),
]) ])
} }
}; }
Selection::Leaf { node_id, leaf_name } => {
frame.render_widget( let node = self.simulation.node(*node_id);
Paragraph::new(body) if let Some(learned) = self.simulation.root_knowledge.node(&node.path)
.block(Block::default().borders(Borders::ALL).title("Inspector")) && let Some(leaf) = learned
.wrap(Wrap { trim: true }), .leaves
area, .iter()
); .find(|leaf| &leaf.leaf_name == leaf_name)
{
let mut lines = vec![
Line::from(format_leaf_ref(&node.path, &leaf.leaf_name)).bold(),
Line::from(
leaf.description
.clone()
.unwrap_or_else(|| "No learned description yet.".to_owned()),
),
Line::from(format!("Node: {}", node.display_path())),
Line::from("Known procedures:"),
];
if leaf.procedures.is_empty() {
lines.push(Line::from("- none learned"));
} else {
lines.extend(leaf.procedures.iter().map(|procedure| {
Line::from(match &procedure.description {
Some(description) => {
format!("- {}: {}", procedure.procedure_id, description)
}
None => format!("- {}", procedure.procedure_id),
})
}));
}
Text::from(lines)
} else {
Text::from(vec![
Line::from(format_leaf_ref(&node.path, leaf_name)).bold(),
Line::from("The root host has not learned this leaf yet."),
])
}
}
}
} }
fn render_trace(&self, frame: &mut Frame<'_>, area: Rect) { fn render_trace(&self, frame: &mut Frame<'_>, area: Rect) {
@@ -211,9 +323,8 @@ impl App {
.map(|hook| { .map(|hook| {
let status = if hook.closed { "closed" } else { "open" }; let status = if hook.closed { "closed" } else { "open" };
ListItem::new(format!( ListItem::new(format!(
"#{} {} -> {} [{}] {}", "{} -> {} [{}] {}",
hook.hook_id, format_hook_ref(&hook.host_path, hook.hook_id),
format_path(&hook.host_path),
format_path(&hook.peer_path), format_path(&hook.peer_path),
status, status,
hook.last_message, hook.last_message,
@@ -230,7 +341,7 @@ impl App {
let help = vec![ let help = vec![
Line::from(self.status.clone()).style(Style::default().add_modifier(Modifier::BOLD)), Line::from(self.status.clone()).style(Style::default().add_modifier(Modifier::BOLD)),
Line::from( Line::from(
"Keys: arrows move selection/scenario | i introspect | e echo leaf | p ping | c chunked | h open chat | d chat data | b chat bye | f invalid peer | s step | a autoplay | q quit", "Keys: arrows move selection/scenario | i introspect | e echo leaf | p ping | c chunked | h open chat | d chat data | b chat bye | f invalid peer | g toggle ground-truth/realistic | m clear deeper memory + realistic | s step | a autoplay | q quit",
), ),
Line::from(format!( Line::from(format!(
"Current selection: {}", "Current selection: {}",
@@ -276,8 +387,21 @@ impl App {
pub(super) fn build_selections(simulation: &Simulation) -> Vec<Selection> { pub(super) fn build_selections(simulation: &Simulation) -> Vec<Selection> {
let mut selections = Vec::new(); let mut selections = Vec::new();
for node in &simulation.tree.nodes { let node_ids: Vec<_> = match simulation.inspector_mode {
InspectorMode::GroundTruth => simulation.tree.nodes.iter().map(|node| node.id).collect(),
InspectorMode::Realistic => simulation
.root_knowledge
.known_paths()
.into_iter()
.filter_map(|path| simulation.tree.find_by_path(&path))
.collect(),
};
for node_id in node_ids {
let node = simulation.node(node_id);
selections.push(Selection::Node(node.id)); selections.push(Selection::Node(node.id));
match simulation.inspector_mode {
InspectorMode::GroundTruth => {
for leaf in &node.leaves { for leaf in &node.leaves {
selections.push(Selection::Leaf { selections.push(Selection::Leaf {
node_id: node.id, node_id: node.id,
@@ -285,5 +409,17 @@ pub(super) fn build_selections(simulation: &Simulation) -> Vec<Selection> {
}); });
} }
} }
InspectorMode::Realistic => {
if let Some(learned) = simulation.root_knowledge.node(&node.path) {
for leaf in &learned.leaves {
selections.push(Selection::Leaf {
node_id: node.id,
leaf_name: leaf.leaf_name.clone(),
});
}
}
}
}
}
selections selections
} }
+10
View File
@@ -182,3 +182,13 @@ pub fn format_path(path: &[String]) -> String {
format!("/{}", path.join("/")) format!("/{}", path.join("/"))
} }
} }
/// Formats a leaf reference using the protocol document's descriptive syntax.
pub fn format_leaf_ref(path: &[String], leaf_name: &str) -> String {
format!("{} {{ leaf: {} }}", format_path(path), leaf_name)
}
/// Formats a hook reference using the protocol document's descriptive syntax.
pub fn format_hook_ref(path: &[String], hook_id: u64) -> String {
format!("{} {{ hook: {} }}", format_path(path), hook_id)
}
+378 -18
View File
@@ -12,14 +12,55 @@ use unshell::protocol::tree::{
ChildRoute, ConnectionState, Endpoint, Ingress, LeafBehavior, LocalEvent, ProtocolEndpoint, ChildRoute, ConnectionState, Endpoint, Ingress, LeafBehavior, LocalEvent, ProtocolEndpoint,
}; };
use unshell::protocol::{ use unshell::protocol::{
CallMessage, DataMessage, FaultMessage, FrameBytes, PacketHeader, PacketType, decode_frame, CallMessage, DataMessage, EndpointIntrospection, FaultMessage, FrameBytes, LeafIntrospection,
PacketHeader, PacketType, decode_frame, deserialize_archived_bytes,
}; };
use crate::model::{ use crate::model::{
DemoTree, EndpointProcedureKind, EndpointProcedureSpec, LeafKind, NodeId, ScenarioDefinition, DemoTree, EndpointProcedureKind, EndpointProcedureSpec, LeafKind, NodeId, ScenarioDefinition,
Selection, format_path, Selection, format_hook_ref, format_leaf_ref, format_path,
}; };
/// Root inspector mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InspectorMode {
GroundTruth,
Realistic,
}
/// Learned procedure metadata stored by the root host.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LearnedProcedure {
pub procedure_id: String,
pub description: Option<String>,
}
/// Learned leaf metadata stored by the root host.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LearnedLeaf {
pub leaf_name: String,
pub description: Option<String>,
pub procedures: Vec<LearnedProcedure>,
}
/// Learned endpoint metadata stored by the root host.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LearnedNode {
pub path: Vec<String>,
pub title: Option<String>,
pub description: Option<String>,
pub direct_child: bool,
pub endpoint_procedures: Vec<LearnedProcedure>,
pub leaves: Vec<LearnedLeaf>,
pub endpoint_introspected: bool,
}
/// Root-host knowledge accumulated from local configuration and observed traffic.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RootKnowledge {
pub nodes: BTreeMap<Vec<String>, LearnedNode>,
}
/// User-facing outcome of a root-originated action. /// User-facing outcome of a root-originated action.
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
pub struct ActionResult { pub struct ActionResult {
@@ -34,6 +75,7 @@ pub struct HookSnapshot {
pub host_path: Vec<String>, pub host_path: Vec<String>,
pub peer_path: Vec<String>, pub peer_path: Vec<String>,
pub procedure_id: String, pub procedure_id: String,
pub target_leaf: Option<String>,
pub closed: bool, pub closed: bool,
pub last_message: String, pub last_message: String,
} }
@@ -121,9 +163,208 @@ pub struct Simulation {
pub trace: VecDeque<TraceEvent>, pub trace: VecDeque<TraceEvent>,
pub recorded_events: Vec<RecordedEvent>, pub recorded_events: Vec<RecordedEvent>,
pub hooks: BTreeMap<u64, HookSnapshot>, pub hooks: BTreeMap<u64, HookSnapshot>,
pub inspector_mode: InspectorMode,
pub root_knowledge: RootKnowledge,
chat_sessions: BTreeMap<u64, ChatSession>, chat_sessions: BTreeMap<u64, ChatSession>,
} }
impl RootKnowledge {
fn new(tree: &DemoTree) -> Self {
let mut knowledge = Self {
nodes: BTreeMap::new(),
};
for node in &tree.nodes {
if node.path.is_empty() || node.path.len() == 1 {
let direct_child = node.path.len() == 1;
let mut learned = LearnedNode {
path: node.path.clone(),
title: Some(node.title.clone()),
description: Some(node.description.clone()),
direct_child,
endpoint_procedures: Vec::new(),
leaves: Vec::new(),
endpoint_introspected: node.path.is_empty(),
};
if node.path.is_empty() {
learned.endpoint_procedures = node
.endpoint_procedures
.iter()
.map(|procedure| LearnedProcedure {
procedure_id: procedure.procedure_id.clone(),
description: Some(procedure.description.clone()),
})
.collect();
learned.leaves = node
.leaves
.iter()
.map(|leaf| LearnedLeaf {
leaf_name: leaf.name.clone(),
description: Some(leaf.description.clone()),
procedures: leaf
.procedures
.iter()
.map(|procedure_id| LearnedProcedure {
procedure_id: procedure_id.clone(),
description: Some(leaf.description.clone()),
})
.collect(),
})
.collect();
}
knowledge.nodes.insert(node.path.clone(), learned);
}
}
knowledge
}
fn ensure_node(&mut self, demo_node: &crate::model::DemoNode) -> &mut LearnedNode {
let direct_child = demo_node.path.len() == 1;
self.nodes
.entry(demo_node.path.clone())
.or_insert_with(|| LearnedNode {
path: demo_node.path.clone(),
title: Some(demo_node.title.clone()),
description: Some(demo_node.description.clone()),
direct_child,
endpoint_procedures: Vec::new(),
leaves: Vec::new(),
endpoint_introspected: false,
})
}
fn remember_endpoint_procedure(
&mut self,
demo_node: &crate::model::DemoNode,
procedure: &EndpointProcedureSpec,
) {
let learned_node = self.ensure_node(demo_node);
push_procedure(
&mut learned_node.endpoint_procedures,
procedure.procedure_id.clone(),
Some(procedure.description.clone()),
);
}
fn remember_leaf_from_spec(
&mut self,
demo_node: &crate::model::DemoNode,
leaf_spec: &crate::model::LeafSpec,
) {
let learned_node = self.ensure_node(demo_node);
let leaf = ensure_leaf(
&mut learned_node.leaves,
leaf_spec.name.clone(),
Some(leaf_spec.description.clone()),
);
for procedure_id in &leaf_spec.procedures {
push_procedure(
&mut leaf.procedures,
procedure_id.clone(),
Some(leaf_spec.description.clone()),
);
}
}
fn remember_endpoint_introspection(
&mut self,
demo_node: &crate::model::DemoNode,
introspection: &EndpointIntrospection,
) {
let learned_node = self.ensure_node(demo_node);
learned_node.endpoint_introspected = true;
for summary in &introspection.leaves {
let description = demo_node
.leaves
.iter()
.find(|leaf| leaf.name == summary.leaf_name)
.map(|leaf| leaf.description.clone());
let leaf = ensure_leaf(
&mut learned_node.leaves,
summary.leaf_name.clone(),
description,
);
for procedure_id in &summary.procedures {
push_procedure(&mut leaf.procedures, procedure_id.clone(), None);
}
}
}
fn remember_leaf_introspection(
&mut self,
demo_node: &crate::model::DemoNode,
introspection: &LeafIntrospection,
) {
let learned_node = self.ensure_node(demo_node);
let description = demo_node
.leaves
.iter()
.find(|leaf| leaf.name == introspection.leaf_name)
.map(|leaf| leaf.description.clone());
let leaf = ensure_leaf(
&mut learned_node.leaves,
introspection.leaf_name.clone(),
description,
);
for procedure_id in &introspection.procedures {
push_procedure(&mut leaf.procedures, procedure_id.clone(), None);
}
}
fn clear_deeper_than_one_hop(&mut self) {
self.nodes.retain(|path, _| path.len() <= 1);
}
pub fn node(&self, path: &[String]) -> Option<&LearnedNode> {
self.nodes.get(path)
}
pub fn known_paths(&self) -> Vec<Vec<String>> {
self.nodes.keys().cloned().collect()
}
}
fn ensure_leaf<'a>(
leaves: &'a mut Vec<LearnedLeaf>,
leaf_name: String,
description: Option<String>,
) -> &'a mut LearnedLeaf {
if let Some(index) = leaves.iter().position(|leaf| leaf.leaf_name == leaf_name) {
if leaves[index].description.is_none() {
leaves[index].description = description;
}
return &mut leaves[index];
}
leaves.push(LearnedLeaf {
leaf_name,
description,
procedures: Vec::new(),
});
leaves.last_mut().expect("just pushed")
}
fn push_procedure(
procedures: &mut Vec<LearnedProcedure>,
procedure_id: String,
description: Option<String>,
) {
if let Some(existing) = procedures
.iter_mut()
.find(|procedure| procedure.procedure_id == procedure_id)
{
if existing.description.is_none() {
existing.description = description;
}
return;
}
procedures.push(LearnedProcedure {
procedure_id,
description,
});
}
impl Simulation { impl Simulation {
/// Creates a fresh simulation from a scenario definition. /// Creates a fresh simulation from a scenario definition.
pub fn new(scenario: ScenarioDefinition) -> Result<Self, SimError> { pub fn new(scenario: ScenarioDefinition) -> Result<Self, SimError> {
@@ -172,6 +413,8 @@ impl Simulation {
}); });
} }
let root_knowledge = RootKnowledge::new(&tree);
Ok(Self { Ok(Self {
scenario, scenario,
tree, tree,
@@ -181,6 +424,8 @@ impl Simulation {
trace: VecDeque::new(), trace: VecDeque::new(),
recorded_events: Vec::new(), recorded_events: Vec::new(),
hooks: BTreeMap::new(), hooks: BTreeMap::new(),
inspector_mode: InspectorMode::GroundTruth,
root_knowledge,
chat_sessions: BTreeMap::new(), chat_sessions: BTreeMap::new(),
}) })
} }
@@ -195,6 +440,25 @@ impl Simulation {
self.tree.node(id) self.tree.node(id)
} }
/// Clears deeper root memory and switches the inspector into realistic mode.
pub fn enable_realistic_mode_with_memory_reset(&mut self) {
self.root_knowledge.clear_deeper_than_one_hop();
self.inspector_mode = InspectorMode::Realistic;
}
/// Toggles the inspector between learned state and ground truth.
pub fn toggle_inspector_mode(&mut self) {
self.inspector_mode = match self.inspector_mode {
InspectorMode::GroundTruth => InspectorMode::Realistic,
InspectorMode::Realistic => InspectorMode::GroundTruth,
};
}
/// Returns whether the inspector is using learned state.
pub fn is_realistic_mode(&self) -> bool {
self.inspector_mode == InspectorMode::Realistic
}
/// Builds and routes an endpoint introspection call from the root. /// Builds and routes an endpoint introspection call from the root.
pub fn call_endpoint_introspection( pub fn call_endpoint_introspection(
&mut self, &mut self,
@@ -215,11 +479,18 @@ impl Simulation {
leaf_name: &str, leaf_name: &str,
) -> Result<ActionResult, SimError> { ) -> Result<ActionResult, SimError> {
let node_path = self.tree.node(node_id).path.clone(); let node_path = self.tree.node(node_id).path.clone();
let node_display = self.tree.node(node_id).display_path();
self.require_leaf(node_id, leaf_name)?; self.require_leaf(node_id, leaf_name)?;
let node = self.tree.node(node_id).clone();
if let Some(leaf_spec) = node.leaves.iter().find(|leaf| leaf.name == leaf_name) {
self.root_knowledge
.remember_leaf_from_spec(&node, leaf_spec);
}
self.dispatch_root_call(node_path, Some(leaf_name.to_owned()), "", Vec::new())?; self.dispatch_root_call(node_path, Some(leaf_name.to_owned()), "", Vec::new())?;
Ok(ActionResult { Ok(ActionResult {
label: format!("Inspect leaf {} on {}", leaf_name, node_display), label: format!(
"Inspect {}",
format_leaf_ref(&self.node(node_id).path, leaf_name)
),
hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id),
}) })
} }
@@ -233,9 +504,18 @@ impl Simulation {
) -> Result<ActionResult, SimError> { ) -> Result<ActionResult, SimError> {
let node_path = self.tree.node(node_id).path.clone(); let node_path = self.tree.node(node_id).path.clone();
let node_display = self.tree.node(node_id).display_path(); let node_display = self.tree.node(node_id).display_path();
let leaf = self.require_leaf(node_id, leaf_name)?; let node = self.tree.node(node_id).clone();
let procedures = self.require_leaf(node_id, leaf_name)?.procedures.clone();
if let Some(leaf_spec) = node
.leaves
.iter()
.find(|known_leaf| known_leaf.name == leaf_name)
{
self.root_knowledge
.remember_leaf_from_spec(&node, leaf_spec);
}
let procedure_id = let procedure_id =
leaf.procedures procedures
.first() .first()
.cloned() .cloned()
.ok_or_else(|| SimError::UnknownProcedure { .ok_or_else(|| SimError::UnknownProcedure {
@@ -249,7 +529,10 @@ impl Simulation {
text.as_bytes().to_vec(), text.as_bytes().to_vec(),
)?; )?;
Ok(ActionResult { Ok(ActionResult {
label: format!("Echo via {leaf_name} on {}", node_display), label: format!(
"Echo via {}",
format_leaf_ref(&self.node(node_id).path, leaf_name)
),
hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id),
}) })
} }
@@ -264,6 +547,15 @@ impl Simulation {
let node_path = self.tree.node(node_id).path.clone(); let node_path = self.tree.node(node_id).path.clone();
let node_display = self.tree.node(node_id).display_path(); let node_display = self.tree.node(node_id).display_path();
self.require_endpoint_procedure(node_id, procedure_id)?; self.require_endpoint_procedure(node_id, procedure_id)?;
let node = self.tree.node(node_id).clone();
if let Some(procedure) = node
.endpoint_procedures
.iter()
.find(|known_procedure| known_procedure.procedure_id == procedure_id)
{
self.root_knowledge
.remember_endpoint_procedure(&node, procedure);
}
self.dispatch_root_call(node_path, None, procedure_id, data)?; self.dispatch_root_call(node_path, None, procedure_id, data)?;
Ok(ActionResult { Ok(ActionResult {
label: format!("Call {procedure_id} on {}", node_display), label: format!("Call {procedure_id} on {}", node_display),
@@ -293,7 +585,10 @@ impl Simulation {
}, },
node_display, node_display,
dst_leaf dst_leaf
.map(|leaf_name| format!(" leaf {leaf_name}")) .map(|leaf_name| format!(
" {}",
format_leaf_ref(&self.node(node_id).path, leaf_name)
))
.unwrap_or_default() .unwrap_or_default()
), ),
hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id), hook_id: self.hooks.last_key_value().map(|(hook_id, _)| *hook_id),
@@ -324,7 +619,10 @@ impl Simulation {
.map_err(|error| SimError::Protocol(error.to_string()))?; .map_err(|error| SimError::Protocol(error.to_string()))?;
self.record_trace( self.record_trace(
self.root_id, self.root_id,
format!("root queued hook data for hook #{hook_id}: {text}"), format!(
"root queued hook data for {}: {text}",
format_hook_ref(self.node(self.root_id).path.as_slice(), hook_id)
),
); );
self.process_local_frame(self.root_id, frame)?; self.process_local_frame(self.root_id, frame)?;
Ok(ActionResult { Ok(ActionResult {
@@ -362,13 +660,17 @@ impl Simulation {
self.record_trace( self.record_trace(
from_node_id, from_node_id,
format!( format!(
"injected invalid peer data toward {} for hook #{hook_id}", "injected invalid peer data toward {} for {}",
format_path(&to_path) format_path(&to_path),
format_hook_ref(self.node(to_node_id).path.as_slice(), hook_id)
), ),
); );
self.process_local_frame(from_node_id, frame)?; self.process_local_frame(from_node_id, frame)?;
Ok(ActionResult { Ok(ActionResult {
label: format!("Inject invalid peer data for hook #{hook_id}"), label: format!(
"Inject invalid peer data for {}",
format_hook_ref(self.node(to_node_id).path.as_slice(), hook_id)
),
hook_id: Some(hook_id), hook_id: Some(hook_id),
}) })
} }
@@ -432,6 +734,7 @@ impl Simulation {
host_path: Vec::new(), host_path: Vec::new(),
peer_path: dst_path.clone(), peer_path: dst_path.clone(),
procedure_id: procedure_id.to_owned(), procedure_id: procedure_id.to_owned(),
target_leaf: dst_leaf.clone(),
closed: false, closed: false,
last_message: format!("created for {}", format_path(&dst_path)), last_message: format!("created for {}", format_path(&dst_path)),
}, },
@@ -448,7 +751,7 @@ impl Simulation {
format_path(&dst_path), format_path(&dst_path),
dst_leaf dst_leaf
.as_ref() .as_ref()
.map(|leaf| format!(" leaf {leaf}")) .map(|leaf| format!(" {}", format_leaf_ref(&dst_path, leaf)))
.unwrap_or_default() .unwrap_or_default()
), ),
); );
@@ -543,8 +846,11 @@ impl Simulation {
self.record_trace( self.record_trace(
node_id, node_id,
format!( format!(
"local Data on hook #{}: {text}", "local Data on {}: {text}",
format_hook_ref(
self.node(node_id).path.as_slice(),
header.hook_id.unwrap_or(0) header.hook_id.unwrap_or(0)
)
), ),
); );
if let Some(hook_id) = header.hook_id { if let Some(hook_id) = header.hook_id {
@@ -558,6 +864,10 @@ impl Simulation {
snapshot.closed = true; snapshot.closed = true;
} }
} }
if node_id == self.root_id {
self.learn_from_root_data(hook_id, &message);
}
} }
if let Some(session) = self if let Some(session) = self
@@ -606,8 +916,11 @@ impl Simulation {
self.record_trace( self.record_trace(
node_id, node_id,
format!( format!(
"local Fault on hook #{}: 0x{:02X}", "local Fault on {}: 0x{:02X}",
header.hook_id.unwrap_or(0), format_hook_ref(
self.node(node_id).path.as_slice(),
header.hook_id.unwrap_or(0)
),
message.fault.0 message.fault.0
), ),
); );
@@ -633,7 +946,7 @@ impl Simulation {
header header
.dst_leaf .dst_leaf
.as_ref() .as_ref()
.map(|leaf| format!("leaf {leaf}")) .map(|leaf| format_leaf_ref(&header.dst_path, leaf))
.unwrap_or_else(|| "endpoint".to_owned()) .unwrap_or_else(|| "endpoint".to_owned())
), ),
); );
@@ -835,8 +1148,55 @@ impl Simulation {
format!("{}: {}", node.display_path(), node.title) format!("{}: {}", node.display_path(), node.title)
} }
Selection::Leaf { node_id, leaf_name } => { Selection::Leaf { node_id, leaf_name } => {
format!("{} leaf {}", self.node(*node_id).display_path(), leaf_name) format_leaf_ref(&self.node(*node_id).path, leaf_name)
} }
} }
} }
fn learn_from_root_data(&mut self, hook_id: u64, message: &DataMessage) {
let Some(snapshot) = self.hooks.get(&hook_id).cloned() else {
return;
};
let Some(node_id) = self.tree.find_by_path(&snapshot.peer_path) else {
return;
};
let demo_node = self.node(node_id).clone();
if snapshot.procedure_id.is_empty() {
if snapshot.target_leaf.is_some() {
if let Ok(introspection) = deserialize_archived_bytes::<
unshell::protocol::introspection::ArchivedLeafIntrospection,
LeafIntrospection,
>(&message.data)
{
self.root_knowledge
.remember_leaf_introspection(&demo_node, &introspection);
}
} else if let Ok(introspection) = deserialize_archived_bytes::<
unshell::protocol::introspection::ArchivedEndpointIntrospection,
EndpointIntrospection,
>(&message.data)
{
self.root_knowledge
.remember_endpoint_introspection(&demo_node, &introspection);
}
return;
}
if let Some(procedure) = demo_node
.endpoint_procedures
.iter()
.find(|procedure| procedure.procedure_id == snapshot.procedure_id)
{
self.root_knowledge
.remember_endpoint_procedure(&demo_node, procedure);
}
if let Some(leaf_name) = &snapshot.target_leaf
&& let Some(leaf_spec) = demo_node.leaves.iter().find(|leaf| &leaf.name == leaf_name)
{
self.root_knowledge
.remember_leaf_from_spec(&demo_node, leaf_spec);
}
}
} }
+41
View File
@@ -0,0 +1,41 @@
use treetest::{model::NodeId, scenarios::built_in_scenarios, sim::Simulation};
#[test]
fn realistic_mode_memory_reset_forgets_deeper_nodes() {
let scenarios = built_in_scenarios();
let mut simulation = Simulation::new(scenarios[2].clone()).expect("scenario should build");
simulation
.call_echo_leaf(NodeId(2), "echo", "learn this nested leaf")
.expect("echo call should start");
simulation.drain().expect("network should drain");
assert!(
simulation
.root_knowledge
.node(&simulation.node(NodeId(2)).path)
.is_some()
);
simulation.enable_realistic_mode_with_memory_reset();
assert!(simulation.is_realistic_mode());
assert!(
simulation
.root_knowledge
.node(&simulation.node(NodeId(2)).path)
.is_none()
);
assert!(
simulation
.root_knowledge
.node(&simulation.node(NodeId(1)).path)
.is_some()
);
assert!(
simulation
.root_knowledge
.node(&simulation.node(NodeId(0)).path)
.is_some()
);
}