Prompt 2

ush-payload/src/main.rs

@@ -7,11 +7,16 @@ use std::io::{Read, Write};
 use std::net::TcpListener;
 use std::sync::mpsc;
 use std::thread;
+use std::time::Duration;
 
-use unshell::tree::{EndpointManager, TreeMessage};
+use serde_json::json;
+use unshell::tree::message::TreeMessage;
+use unshell::tree::protocols::{ProtocolConfig, ProtocolStack};
+use unshell::tree::tcp::{TcpClient, TcpServer};
+use unshell::tree::{ComponentRegistry, EndpointManager, TreeElement};
 
 fn main() {
-    println!("=== TCP Chain Test Harness ===\n");
+    println!("=== Tree Protocol Test Harness ===\n");
 
     // Test 1: Local TCP Server-Client loopback
     test_tcp_loopback();
@@ -19,9 +24,21 @@ fn main() {
     // Test 2: Tree message routing
     test_tree_message();
 
-    // Test 3: TreeMessage serialization
+    // Test 3: TreeMessage serialization (new API)
     test_message_serialization();
 
+    // Test 4: TCP Server with RPC
+    test_tcp_server();
+
+    // Test 5: TCP Client with RPC
+    test_tcp_client();
+
+    // Test 6: Protocol stacking
+    test_protocol_stack();
+
+    // Test 7: Component registry
+    test_component_registry();
+
     println!("\n=== All tests complete ===");
 }
 
@@ -29,7 +46,6 @@ fn main() {
 fn test_tcp_loopback() {
     println!("[Test 1] TCP Loopback Test");
 
-    // Start a TCP server in a thread
     let (tx, rx) = mpsc::channel();
 
     let server_thread = thread::spawn(move || {
@@ -37,24 +53,25 @@ fn test_tcp_loopback() {
         let addr = listener.local_addr().unwrap();
         tx.send(addr.port()).unwrap();
 
-        for stream in listener.incoming() {
-            if let Ok(mut stream) = stream {
-                let mut buf = [0u8; 1024];
-                if let Ok(n) = stream.read(&mut buf) {
-                    let response = b"Echo: ";
-                    let _ = stream.write(response);
-                    let _ = stream.write(&buf[..n]);
-                    let _ = stream.flush();
-                }
+        // Accept one connection only
+        if let Ok((mut stream, _)) = listener.accept() {
+            let mut buf = [0u8; 1024];
+            if let Ok(n) = stream.read(&mut buf) {
+                let response = b"Echo: ";
+                let _ = stream.write(response);
+                let _ = stream.write(&buf[..n]);
+                let _ = stream.flush();
             }
         }
     });
 
     let port = rx.recv().unwrap();
 
-    // Connect client
     let mut stream =
         std::net::TcpStream::connect(format!("127.0.0.1:{}", port)).expect("Failed to connect");
+    stream
+        .set_read_timeout(Some(Duration::from_millis(1000)))
+        .unwrap();
 
     let msg = b"Hello from client!";
     stream.write(msg).expect("Failed to write");
@@ -76,7 +93,6 @@ fn test_tree_message() {
 
     let mut endpoint = EndpointManager::new("endpoint-1");
 
-    // Test GetChildren
     let response = endpoint
         .branch_mut()
         .send_message(serde_json::Value::Null, serde_json::json!("GetChildren"));
@@ -84,13 +100,11 @@ fn test_tree_message() {
     let children = response.as_object().unwrap();
     println!("  Children: {:?}", children.keys().collect::<Vec<_>>());
 
-    // Test ID access
     let response = endpoint
         .branch_mut()
         .send_message(serde_json::json!("id"), serde_json::Value::Null);
     println!("  Endpoint ID: {:?}", response);
 
-    // Test logs queue
     let sender = endpoint.logs_sender().clone();
     sender.send(serde_json::json!("Test log entry")).unwrap();
 
@@ -102,24 +116,168 @@ fn test_tree_message() {
     println!("  ✓ Tree message test passed\n");
 }
 
-/// Test TreeMessage serialization
+/// Test TreeMessage serialization (new API)
 fn test_message_serialization() {
     println!("[Test 3] TreeMessage Serialization");
 
-    let msg = TreeMessage::new_req(
-        "msg-1",
-        vec!["endpoint1".to_string(), "shell".to_string()],
-        "Get",
-    );
+    // Test new API
+    let msg = TreeMessage::new("query")
+        .to_target(["endpoint1", "shell"])
+        .with_payload(json!({"key": "value"}));
 
     let json = serde_json::to_string_pretty(&msg).unwrap();
     println!("  Message: {}", json);
 
-    // Test response message
-    let resp = TreeMessage::new_resp("resp-1", "msg-1", serde_json::json!("result data"));
+    // Test response
+    let resp =
+        TreeMessage::new("response").with_payload(json!({"success": true, "result": "data"}));
 
     let json = serde_json::to_string_pretty(&resp).unwrap();
     println!("  Response: {}", json);
 
+    // Test RPC call
+    let rpc_msg = TreeMessage::new("rpc.call")
+        .to_target(["components", "tcp-client"])
+        .with_payload(json!({
+            "method": "connect",
+            "params": {"address": "127.0.0.1", "port": 8080}
+        }));
+
+    let json = serde_json::to_string_pretty(&rpc_msg).unwrap();
+    println!("  RPC Call: {}", json);
+
     println!("  ✓ Message serialization test passed\n");
 }
+
+/// Test TCP Server with RPC
+fn test_tcp_server() {
+    println!("[Test 4] TCP Server with RPC");
+
+    let mut server = TcpServer::new("test-server");
+
+    // Configure
+    let response = server.send_message(
+        json!(null),
+        json!({
+            "method": "status"
+        }),
+    );
+    println!("  Initial status: {:?}", response);
+
+    // Try to start listening
+    let response = server.send_message(
+        json!(null),
+        json!({
+            "method": "listen",
+            "params": {
+                "bind_address": "127.0.0.1",
+                "port": 0
+            }
+        }),
+    );
+    println!("  Listen result: {:?}", response);
+
+    // Get status after listening
+    let response = server.send_message(json!(null), json!({"method": "status"}));
+    println!("  Status after listen: {:?}", response);
+
+    println!("  ✓ TCP Server test passed\n");
+}
+
+/// Test TCP Client with RPC
+fn test_tcp_client() {
+    println!("[Test 5] TCP Client with RPC");
+
+    let mut client = TcpClient::new("test-client");
+
+    // Get initial status
+    let response = client.send_message(json!(null), json!({"method": "status"}));
+    println!("  Initial status: {:?}", response);
+
+    // Try to connect (will fail since no server, but tests the RPC)
+    let response = client.send_message(
+        json!(null),
+        json!({
+            "method": "connect",
+            "params": {"address": "127.0.0.1", "port": 65432}
+        }),
+    );
+    println!("  Connect result: {:?}", response);
+
+    // Get status after connect attempt
+    let response = client.send_message(json!(null), json!({"method": "status"}));
+    println!("  Status after connect: {:?}", response);
+
+    println!("  ✓ TCP Client test passed\n");
+}
+
+/// Test Protocol stacking
+fn test_protocol_stack() {
+    println!("[Test 6] Protocol Stacking");
+
+    // Create a stack: base64 -> tcp
+    let mut stack = ProtocolStack::new();
+    stack
+        .push(&ProtocolConfig::Base64(Default::default()))
+        .unwrap();
+    stack
+        .push(&ProtocolConfig::Tcp(Default::default()))
+        .unwrap();
+
+    println!("  Stack protocols: {:?}", stack.to_configs());
+
+    // Test encoding/decoding
+    let test_data = b"Hello, World!";
+    let encoded = stack.encode(test_data).unwrap();
+    println!(
+        "  Encoded ({} bytes): {:?}",
+        encoded.len(),
+        String::from_utf8_lossy(&encoded)
+    );
+
+    let decoded = stack.decode(&encoded).unwrap();
+    println!("  Decoded: {:?}", String::from_utf8_lossy(&decoded));
+
+    // Test with HTTP
+    let mut http_stack = ProtocolStack::new();
+    http_stack
+        .push(&ProtocolConfig::Base64(Default::default()))
+        .unwrap();
+    http_stack
+        .push(&ProtocolConfig::Http(Default::default()))
+        .unwrap();
+
+    let test_data = b"test message";
+    let encoded = http_stack.encode(test_data).unwrap();
+    println!(
+        "  HTTP+Base64 encoded: {:?}",
+        String::from_utf8_lossy(&encoded).lines().next()
+    );
+
+    println!("  ✓ Protocol stacking test passed\n");
+}
+
+/// Test Component Registry
+fn test_component_registry() {
+    println!("[Test 7] Component Registry");
+
+    let mut registry = ComponentRegistry::new();
+
+    // Create and register a TCP client component
+    let client = Box::new(TcpClient::new("tcp-client-1"));
+    registry.register(client).unwrap();
+
+    // List components
+    let list = registry.list();
+    println!("  Registered components: {:?}", list);
+
+    // Send RPC to component
+    let result = registry.send_to_component("tcp-client-1", json!({"method": "status"}));
+    println!("  Component status: {:?}", result);
+
+    // Send RPC via path
+    let result = registry.send_message(json!("rpc.tcp-client-1"), json!({"method": "status"}));
+    println!("  Via RPC path: {:?}", result);
+
+    println!("  ✓ Component registry test passed\n");
+}