fix: markdown rendering, continuous polling, and attack chain graph loading

- Add markdown-to-HTML renderer for assistant messages (headers, bold,
  code blocks, lists, inline code)
- Fix polling to continuously loop while session is running using
  poll_gen signal
- Fix attack chain graph loading with spawn delay for DOM readiness
- Default attack chain tab to list view (more reliable)
- Render tool_result role messages as tool indicators

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Sharang Parnerkar
2026-03-11 21:59:36 +01:00
parent 30301a12b5
commit 0428cba2b8

View File

@@ -8,6 +8,157 @@ use crate::infrastructure::pentest::{
fetch_pentest_session, send_pentest_message,
};
/// Simple markdown-to-HTML converter for assistant messages.
/// Handles headers, bold, italic, code blocks, inline code, and lists.
fn markdown_to_html(input: &str) -> String {
let mut html = String::new();
let mut in_code_block = false;
let mut in_list = false;
for line in input.lines() {
if line.starts_with("```") {
if in_code_block {
html.push_str("</code></pre>");
in_code_block = false;
} else {
if in_list {
html.push_str("</ul>");
in_list = false;
}
html.push_str("<pre style=\"background:var(--bg-primary);padding:10px;border-radius:6px;overflow-x:auto;font-size:0.8rem;margin:6px 0;\"><code>");
in_code_block = true;
}
continue;
}
if in_code_block {
// Escape HTML inside code blocks
let escaped = line
.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;");
html.push_str(&escaped);
html.push('\n');
continue;
}
let trimmed = line.trim();
// Blank line — close list if open
if trimmed.is_empty() {
if in_list {
html.push_str("</ul>");
in_list = false;
}
html.push_str("<br/>");
continue;
}
// Lists
if trimmed.starts_with("- ") || trimmed.starts_with("* ") {
if !in_list {
html.push_str("<ul style=\"margin:4px 0;padding-left:20px;\">");
in_list = true;
}
let content = inline_format(&trimmed[2..]);
html.push_str(&format!("<li>{content}</li>"));
continue;
}
// Numbered lists
if trimmed.len() > 2 {
let mut chars = trimmed.chars();
let first = chars.next();
let second = chars.next();
if first.map(|c| c.is_ascii_digit()).unwrap_or(false)
&& (second == Some('.') || second == Some(')'))
{
let rest = &trimmed[2..].trim_start();
if !in_list {
html.push_str("<ul style=\"margin:4px 0;padding-left:20px;\">");
in_list = true;
}
let content = inline_format(rest);
html.push_str(&format!("<li>{content}</li>"));
continue;
}
}
// Close list if we're no longer in one
if in_list {
html.push_str("</ul>");
in_list = false;
}
// Headers
if trimmed.starts_with("### ") {
let content = inline_format(&trimmed[4..]);
html.push_str(&format!(
"<h4 style=\"margin:8px 0 4px;font-size:0.9rem;\">{content}</h4>"
));
} else if trimmed.starts_with("## ") {
let content = inline_format(&trimmed[3..]);
html.push_str(&format!(
"<h3 style=\"margin:10px 0 4px;font-size:0.95rem;\">{content}</h3>"
));
} else if trimmed.starts_with("# ") {
let content = inline_format(&trimmed[2..]);
html.push_str(&format!(
"<h2 style=\"margin:12px 0 6px;font-size:1rem;\">{content}</h2>"
));
} else {
let content = inline_format(trimmed);
html.push_str(&format!("<p style=\"margin:2px 0;\">{content}</p>"));
}
}
if in_list {
html.push_str("</ul>");
}
if in_code_block {
html.push_str("</code></pre>");
}
html
}
/// Handle inline formatting: bold, italic, inline code
fn inline_format(text: &str) -> String {
let mut result = text
.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;");
// Inline code (backticks)
while let Some(start) = result.find('`') {
if let Some(end) = result[start + 1..].find('`') {
let code_content = &result[start + 1..start + 1 + end].to_string();
let replacement = format!(
"<code style=\"background:var(--bg-primary);padding:1px 4px;border-radius:3px;font-size:0.85em;\">{code_content}</code>"
);
result = format!("{}{}{}", &result[..start], replacement, &result[start + 2 + end..]);
} else {
break;
}
}
// Bold (**text**)
while let Some(start) = result.find("**") {
if let Some(end) = result[start + 2..].find("**") {
let bold_content = &result[start + 2..start + 2 + end].to_string();
result = format!(
"{}<strong>{bold_content}</strong>{}",
&result[..start],
&result[start + 4 + end..]
);
} else {
break;
}
}
result
}
#[component]
pub fn PentestSessionPage(session_id: String) -> Element {
let sid = session_id.clone();
@@ -35,10 +186,11 @@ pub fn PentestSessionPage(session_id: String) -> Element {
let mut input_text = use_signal(String::new);
let mut sending = use_signal(|| false);
let mut right_tab = use_signal(|| "findings".to_string());
let mut chain_view = use_signal(|| "graph".to_string()); // "graph" or "list"
let mut chain_view = use_signal(|| "list".to_string());
let mut exporting = use_signal(|| false);
let mut poll_gen = use_signal(|| 0u32); // incremented to trigger re-poll
// Auto-poll messages every 3s when session is running
// Extract session status
let session_status = {
let s = session.read();
match &*s {
@@ -54,10 +206,10 @@ pub fn PentestSessionPage(session_id: String) -> Element {
let is_running = session_status == "running";
let sid_for_poll = session_id.clone();
// Continuous polling: re-fetch all data every 3s while running
use_effect(move || {
let _gen = *poll_gen.read(); // subscribe to changes
if is_running {
let _sid = sid_for_poll.clone();
spawn(async move {
#[cfg(feature = "web")]
gloo_timers::future::TimeoutFuture::new(3_000).await;
@@ -67,24 +219,36 @@ pub fn PentestSessionPage(session_id: String) -> Element {
findings.restart();
attack_chain.restart();
session.restart();
// Bump generation to trigger the next poll cycle
let next = poll_gen.peek().wrapping_add(1);
poll_gen.set(next);
});
}
});
// Load attack chain into vis-network when data changes and graph tab is active
let chain_data_for_viz = attack_chain.read().clone();
let current_tab = right_tab.read().clone();
let current_chain_view = chain_view.read().clone();
// Load attack chain into vis-network when graph tab is active and data is available
// Use a separate effect that reads the reactive resources directly
use_effect(move || {
if current_tab == "chain" && current_chain_view == "graph" {
if let Some(Some(data)) = &chain_data_for_viz {
let tab = right_tab.read().clone();
let view = chain_view.read().clone();
let chain = attack_chain.read().clone();
if tab == "chain" && view == "graph" {
if let Some(Some(data)) = &chain {
if !data.data.is_empty() {
let nodes_json =
serde_json::to_string(&data.data).unwrap_or_else(|_| "[]".to_string());
let js = format!(
r#"if (window.__loadAttackChain) {{ window.__loadAttackChain({nodes_json}); }}"#
);
document::eval(&js);
// Small delay to ensure the DOM container exists
spawn(async move {
#[cfg(feature = "web")]
gloo_timers::future::TimeoutFuture::new(100).await;
#[cfg(not(feature = "web"))]
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let js = format!(
r#"if (window.__loadAttackChain) {{ window.__loadAttackChain({nodes_json}); }}"#
);
document::eval(&js);
});
}
}
}
@@ -109,7 +273,7 @@ pub fn PentestSessionPage(session_id: String) -> Element {
let mut do_send_click = do_send.clone();
// Export handler
// Export handlers
let sid_for_export = session_id.clone();
let do_export_md = move |_| {
let sid = sid_for_export.clone();
@@ -117,7 +281,6 @@ pub fn PentestSessionPage(session_id: String) -> Element {
spawn(async move {
match export_pentest_report(sid.clone(), "markdown".to_string()).await {
Ok(content) => {
// Trigger download via JS
let escaped = content
.replace('\\', "\\\\")
.replace('`', "\\`")
@@ -256,7 +419,6 @@ pub fn PentestSessionPage(session_id: String) -> Element {
}
}
div { style: "display: flex; gap: 8px; align-items: center;",
// Export buttons
div { style: "display: flex; gap: 4px;",
button {
class: "btn btn-ghost",
@@ -318,7 +480,7 @@ pub fn PentestSessionPage(session_id: String) -> Element {
let tool_name = msg.get("tool_name").and_then(|v| v.as_str()).unwrap_or("").to_string();
let tool_status = msg.get("tool_status").and_then(|v| v.as_str()).unwrap_or("").to_string();
if msg_type == "tool_call" || msg_type == "tool_result" {
if msg_type == "tool_call" || msg_type == "tool_result" || role == "tool_result" {
let tool_icon_style = match tool_status.as_str() {
"success" => "color: #16a34a;",
"error" => "color: #dc2626;",
@@ -358,6 +520,8 @@ pub fn PentestSessionPage(session_id: String) -> Element {
}
}
} else {
// Assistant message — render markdown
let rendered_html = markdown_to_html(&content);
rsx! {
div {
key: "{i}",
@@ -367,8 +531,8 @@ pub fn PentestSessionPage(session_id: String) -> Element {
Icon { icon: BsCpu, width: 14, height: 14 }
}
div {
style: "max-width: 80%; padding: 10px 14px; background: var(--bg-tertiary); border-radius: 12px 12px 12px 2px; font-size: 0.9rem; line-height: 1.5; white-space: pre-wrap;",
"{content}"
style: "max-width: 80%; padding: 10px 14px; background: var(--bg-tertiary); border-radius: 12px 12px 12px 2px; font-size: 0.9rem; line-height: 1.5;",
dangerous_inner_html: "{rendered_html}",
}
}
}