From cc6ae7717ccd11ebb33595903e9cca064d6d8c92 Mon Sep 17 00:00:00 2001
From: Sharang Parnerkar <parnerkarsharang@gmail.com>
Date: Wed, 11 Mar 2026 22:12:14 +0100
Subject: [PATCH] fix: attack chain node linking and disable input while
 pentest runs

Link attack chain nodes to previous iteration's nodes via parent_node_ids
so the DAG graph shows proper hierarchy instead of flat dots. Disable the
chat input while a pentest session is running since messages have no effect.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 compliance-agent/src/pentest/orchestrator.rs  |  9 ++++
 .../src/pages/pentest_session.rs              | 47 +++++++++++--------
 2 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/compliance-agent/src/pentest/orchestrator.rs b/compliance-agent/src/pentest/orchestrator.rs
index 4fd9e2f..482902f 100644
--- a/compliance-agent/src/pentest/orchestrator.rs
+++ b/compliance-agent/src/pentest/orchestrator.rs
@@ -172,6 +172,7 @@ impl PentestOrchestrator {
         let mut total_findings = 0u32;
         let mut total_tool_calls = 0u32;
         let mut total_successes = 0u32;
+        let mut prev_node_ids: Vec<String> = Vec::new();
 
         for _iteration in 0..max_iterations {
             let response = self
@@ -233,6 +234,8 @@ impl PentestOrchestrator {
                         tool_call_id: None,
                     });
 
+                    let mut current_batch_node_ids: Vec<String> = Vec::new();
+
                     for tc in &tool_calls {
                         total_tool_calls += 1;
                         let node_id = uuid::Uuid::new_v4().to_string();
@@ -244,9 +247,12 @@ impl PentestOrchestrator {
                             tc.arguments.clone(),
                             String::new(),
                         );
+                        // Link to previous iteration's nodes
+                        node.parent_node_ids = prev_node_ids.clone();
                         node.status = AttackNodeStatus::Running;
                         node.started_at = Some(chrono::Utc::now());
                         let _ = self.db.attack_chain_nodes().insert_one(&node).await;
+                        current_batch_node_ids.push(node_id.clone());
 
                         let _ = self.event_tx.send(PentestEvent::ToolStart {
                             node_id: node_id.clone(),
@@ -337,6 +343,9 @@ impl PentestOrchestrator {
                         });
                     }
 
+                    // Advance parent links so next iteration's nodes connect to this batch
+                    prev_node_ids = current_batch_node_ids;
+
                     if let Some(sid) = session.id {
                         let _ = self
                             .db
diff --git a/compliance-dashboard/src/pages/pentest_session.rs b/compliance-dashboard/src/pages/pentest_session.rs
index f8f13fa..f8a152c 100644
--- a/compliance-dashboard/src/pages/pentest_session.rs
+++ b/compliance-dashboard/src/pages/pentest_session.rs
@@ -560,27 +560,34 @@ pub fn PentestSessionPage(session_id: String) -> Element {
                     }
                 }
 
-                // Input area
-                div { style: "flex-shrink: 0; padding: 12px; border-top: 1px solid var(--border-color); display: flex; gap: 8px;",
-                    textarea {
-                        class: "chat-input",
-                        style: "flex: 1;",
-                        placeholder: "Guide the pentest agent...",
-                        value: "{input_text}",
-                        oninput: move |e| input_text.set(e.value()),
-                        onkeydown: move |e: Event<KeyboardData>| {
-                            if e.key() == Key::Enter && !e.modifiers().shift() {
-                                e.prevent_default();
-                                do_send();
-                            }
-                        },
+                // Input area — disabled while pentest is running (messages have no effect)
+                if is_running {
+                    div { style: "flex-shrink: 0; padding: 12px; border-top: 1px solid var(--border-color); text-align: center; color: var(--text-secondary); font-size: 0.85rem;",
+                        Icon { icon: BsPlayCircle, width: 14, height: 14 }
+                        " Pentest is running — input disabled until complete"
                     }
-                    button {
-                        class: "btn btn-primary",
-                        style: "align-self: flex-end;",
-                        disabled: *sending.read(),
-                        onclick: move |_| do_send_click(),
-                        "Send"
+                } else {
+                    div { style: "flex-shrink: 0; padding: 12px; border-top: 1px solid var(--border-color); display: flex; gap: 8px;",
+                        textarea {
+                            class: "chat-input",
+                            style: "flex: 1;",
+                            placeholder: "Guide the pentest agent...",
+                            value: "{input_text}",
+                            oninput: move |e| input_text.set(e.value()),
+                            onkeydown: move |e: Event<KeyboardData>| {
+                                if e.key() == Key::Enter && !e.modifiers().shift() {
+                                    e.prevent_default();
+                                    do_send();
+                                }
+                            },
+                        }
+                        button {
+                            class: "btn btn-primary",
+                            style: "align-self: flex-end;",
+                            disabled: *sending.read(),
+                            onclick: move |_| do_send_click(),
+                            "Send"
+                        }
                     }
                 }
             }