From 06715d11bd8b6551e0e79d14cfb305b7635b0e48 Mon Sep 17 00:00:00 2001 From: CyberSecurityUP Date: Sun, 18 Oct 2026 19:19:48 -0300 Subject: [PATCH] feat(harness): add execution resilience module with error classification, backoff, and crash salvage - Add resilience.rs: transient error classification (7 categories), exponential backoff (5 retries, 2s-90s cap), tool-call output limiter, context compaction, and crash-state salvage - Upgrade pool.rs retry loop to use classified errors and exponential backoff instead of fixed 3-retry quadratic delay - Wire crash salvage into all 8 exploit error paths so partial findings survive agent failures - Make extract_findings pub(crate) for cross-module access --- neurosploit-rs/crates/harness/src/lib.rs | 1 + neurosploit-rs/crates/harness/src/pipeline.rs | 96 ++++- neurosploit-rs/crates/harness/src/pool.rs | 31 +- .../crates/harness/src/resilience.rs | 407 ++++++++++++++++++ 4 files changed, 510 insertions(+), 25 deletions(-) create mode 100644 neurosploit-rs/crates/harness/src/resilience.rs diff --git a/neurosploit-rs/crates/harness/src/lib.rs b/neurosploit-rs/crates/harness/src/lib.rs index 4010b00..110c102 100644 --- a/neurosploit-rs/crates/harness/src/lib.rs +++ b/neurosploit-rs/crates/harness/src/lib.rs @@ -47,6 +47,7 @@ pub mod pool; pub mod probe; pub mod replay; pub mod report; +pub mod resilience; pub mod rl; pub mod sandbox; pub mod sarif; diff --git a/neurosploit-rs/crates/harness/src/pipeline.rs b/neurosploit-rs/crates/harness/src/pipeline.rs index 721476d..f88f400 100644 --- a/neurosploit-rs/crates/harness/src/pipeline.rs +++ b/neurosploit-rs/crates/harness/src/pipeline.rs @@ -1445,7 +1445,12 @@ pub async fn run(mut cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sender } else { let _ = txc.send(format!("exploit {} failed: {e}", ag.name)).await; } - (ag.name.clone(), format!("ERROR: {e}"), vec![]) + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { + let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; + } + (ag.name.clone(), error_text, salvaged) } } } @@ -1454,9 +1459,16 @@ pub async fn run(mut cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sender .collect() .await; + // Merge any previously salvaged findings from prior crashed runs. + let mut prior_salvaged = if let Some(ref wd) = cfg.workdir { + crate::resilience::load_salvaged(wd) + } else { Vec::new() }; + crate::traces::record_all(cfg.workdir.as_deref(), &raw); let transcript = transcript_of(&raw); - let candidates = dedup_findings(raw.iter().flat_map(|(_, _, f)| f.clone()).collect()); + let mut all_findings: Vec = raw.iter().flat_map(|(_, _, f)| f.clone()).collect(); + all_findings.append(&mut prior_salvaged); + let candidates = dedup_findings(all_findings); let _ = tx.send(format!("{} candidate finding(s) (deduped) — validating by {}-model vote", candidates.len(), cfg.vote_n)).await; if pool.candidates.len() == 1 && effective_vote_n(&cfg) <= 1 { let _ = tx.send("⚠ single-model panel with vote_n=1 — validation is weaker (same model validates its own findings). Consider --vote-n 2 or adding a second model for cross-validation.".into()).await; @@ -1528,10 +1540,12 @@ pub async fn run_whitebox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: S return RunOutput { target: cfg.target.clone(), workdir: cfg.workdir.clone().unwrap_or_default(), findings: vec![], agents_ran: selected.iter().map(|a| a.name.clone()).collect(), candidates: 0, recon: String::new(), artifacts, denied: None }; } + let wb_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let ctx = context.clone(); let txc = tx.clone(); + let wd = wb_workdir.clone(); async move { let user = format!( "{}\n\nSOURCE CODE TO REVIEW:\n```\n{}\n```\n\nReply ONLY with a JSON array of findings (may be empty []). \ @@ -1541,9 +1555,6 @@ pub async fn run_whitebox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: S ag.user.replace("{target}", "the provided repository").replace("{recon_json}", "{}"), ctx ); - // Prepend the white-box doctrine so code agents stay in static - // source-review mode and never hallucinate live/black-box actions. - // In research mode, add the novelty/patch-diff hunting doctrine. let sys = if research { format!("{}{}{}", WHITEBOX_DOCTRINE, WHITEBOX_RESEARCH_DOCTRINE, ag.system) } else { @@ -1557,7 +1568,12 @@ pub async fn run_whitebox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: S } Err(e) => { let _ = txc.send(format!("analyze {} failed: {e}", ag.name)).await; - (ag.name.clone(), format!("ERROR: {e}"), vec![]) + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { + let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; + } + (ag.name.clone(), error_text, salvaged) } } } @@ -1682,6 +1698,7 @@ pub async fn run_greybox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Se let ops = engagement_ops(&cfg); let recon_ctx: String = recon.chars().take(3000).collect(); let leads_ctx = code_leads.clone(); + let gb_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let target = target.clone(); @@ -1690,6 +1707,7 @@ pub async fn run_greybox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Se let ops = ops.clone(); let leads = leads_ctx.clone(); let txc = tx.clone(); + let wd = gb_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); @@ -1711,8 +1729,15 @@ pub async fn run_greybox(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Se Ok((m, text)) => { let f = extract_findings(&text, &ag.name); let _ = txc.send(format!("exploit {} via {} → {} candidate(s)", ag.name, m.label(), f.len())).await; (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("exploit {} failed: {e}", ag.name)).await; - (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("exploit {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { + let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; + } + (ag.name.clone(), error_text, salvaged) + } } } }) @@ -3633,7 +3658,7 @@ fn findings_items(val: serde_json::Value) -> Vec { /// runs (`JSON parse failed` and `no JSON array/object found`) now collapse into /// one honest outcome: either we recover a value, or the reply held no JSON. /// The recovered value's *shape* is then normalised by [`findings_items`]. -fn extract_findings(text: &str, agent: &str) -> Vec { +pub(crate) fn extract_findings(text: &str, agent: &str) -> Vec { let val = match crate::json_extract::parse_reply(text) { Some(v) => v, None => { @@ -4247,12 +4272,14 @@ pub async fn run_container(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: let verbose = cfg.verbose; let directives = operator_directives(&cfg); let recon_ctx: String = recon.chars().take(3000).collect(); + let ct_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let target = target.clone(); let recon = recon_ctx.clone(); let directives = directives.clone(); let txc = tx.clone(); + let wd = ct_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); } if verbose { let _ = txc.send(format!(" launching agent: {} ({})", ag.name, ag.title.replace(" Agent", ""))).await; } @@ -4268,7 +4295,15 @@ pub async fn run_container(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: for c in &f { if let Ok(j) = serde_json::to_string(c) { let _ = txc.send(format!("finding_json: {j}")).await; } } (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("scan {} failed: {e}", ag.name)).await; (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("scan {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { + let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; + } + (ag.name.clone(), error_text, salvaged) + } } } }) @@ -4338,12 +4373,14 @@ pub async fn run_mobile(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sen let verbose = cfg.verbose; let directives = operator_directives(&cfg); let recon_ctx: String = recon.chars().take(3000).collect(); + let mob_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let target = target.clone(); let recon = recon_ctx.clone(); let directives = directives.clone(); let txc = tx.clone(); + let wd = mob_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); } if verbose { let _ = txc.send(format!(" launching agent: {} ({})", ag.name, ag.title.replace(" Agent", ""))).await; } @@ -4359,7 +4396,13 @@ pub async fn run_mobile(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sen for c in &f { if let Ok(j) = serde_json::to_string(c) { let _ = txc.send(format!("finding_json: {j}")).await; } } (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("test {} failed: {e}", ag.name)).await; (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("test {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; } + (ag.name.clone(), error_text, salvaged) + } } } }) @@ -4428,12 +4471,14 @@ pub async fn run_host(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sende let verbose = cfg.verbose; let directives = operator_directives(&cfg); let recon_ctx: String = recon.chars().take(3000).collect(); + let host_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let target = target.clone(); let recon = recon_ctx.clone(); let directives = directives.clone(); let txc = tx.clone(); + let wd = host_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); } if verbose { @@ -4454,8 +4499,13 @@ pub async fn run_host(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sende } (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("test {} failed: {e}", ag.name)).await; - (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("test {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; } + (ag.name.clone(), error_text, salvaged) + } } } }) @@ -4948,9 +4998,11 @@ pub async fn run_ai(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sender< let target = cfg.target.clone(); let directives = operator_directives(&cfg); let recon_ctx: String = recon.chars().take(3500).collect(); + let ai_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(selected.iter().cloned()) .map(|ag| { let (target, recon, directives, txc) = (target.clone(), recon_ctx.clone(), directives.clone(), tx.clone()); + let wd = ai_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); } let _ = txc.send(format!(" ▶ AI test: {} ({})", ag.name, ag.title.replace(" Agent", ""))).await; @@ -4969,7 +5021,13 @@ pub async fn run_ai(cfg: RunConfig, lib: &Library, pool: &ModelPool, tx: Sender< } (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("ai {} failed: {e}", ag.name)).await; (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("ai {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; } + (ag.name.clone(), error_text, salvaged) + } } } }) @@ -5008,9 +5066,11 @@ pub async fn run_skills_audit(cfg: RunConfig, lib: &Library, pool: &ModelPool, t return finish(cfg, lib, pool, "{}".into(), String::new(), vec![], agents, &mut rl, crate::grounding::GroundMode::Symbolic, String::new(), tx).await; } let directives = operator_directives(&cfg); + let sk_workdir = cfg.workdir.clone(); let raw: Vec<(String, String, Vec)> = stream::iter(agents.iter().cloned()) .map(|ag| { let (ctx, dir, txc) = (context.clone(), directives.clone(), tx.clone()); + let wd = sk_workdir.clone(); async move { if pool.stop_exploiting() { return (ag.name.clone(), String::new(), vec![]); } let _ = txc.send(format!(" ▶ skill audit: {}", ag.name)).await; @@ -5026,7 +5086,13 @@ pub async fn run_skills_audit(cfg: RunConfig, lib: &Library, pool: &ModelPool, t for c in &f { if let Ok(j) = serde_json::to_string(c) { let _ = txc.send(format!("finding_json: {j}")).await; } } (ag.name.clone(), text, f) } - Err(e) => { let _ = txc.send(format!("skill {} failed: {e}", ag.name)).await; (ag.name.clone(), format!("ERROR: {e}"), vec![]) } + Err(e) => { + let _ = txc.send(format!("skill {} failed: {e}", ag.name)).await; + let error_text = format!("ERROR: {e}"); + let salvaged = crate::resilience::salvage_findings(&ag.name, &error_text, wd.as_deref()); + if !salvaged.is_empty() { let _ = txc.send(format!("🔧 salvaged {} finding(s) from crashed agent {}", salvaged.len(), ag.name)).await; } + (ag.name.clone(), error_text, salvaged) + } } } }) diff --git a/neurosploit-rs/crates/harness/src/pool.rs b/neurosploit-rs/crates/harness/src/pool.rs index e6f28fd..f9ce1a8 100644 --- a/neurosploit-rs/crates/harness/src/pool.rs +++ b/neurosploit-rs/crates/harness/src/pool.rs @@ -276,18 +276,20 @@ impl ModelPool { /// HTTP API, with a short retry/backoff. `label` (e.g. the agent name) tags /// the streamed activity so each command/tool is attributable. async fn one(&self, label: &str, m: &ModelRef, system: &str, user: &str) -> Result { + use crate::resilience::{classify_error, ErrorClass, MAX_TRANSIENT_RETRIES, retry_delay}; + if self.is_cancelled() { return Err(anyhow!("cancelled")); } let use_cli = self.subscription && cli_binary_for(&m.provider).is_some(); let progress = self.progress(); let mut last = anyhow::anyhow!("no attempt"); - for attempt in 0..3u64 { + for attempt in 0..MAX_TRANSIENT_RETRIES { if self.is_cancelled() { return Err(anyhow!("cancelled")); } if attempt > 0 { - tokio::time::sleep(std::time::Duration::from_millis(1500 * attempt * attempt.max(1))).await; + tokio::time::sleep(retry_delay(attempt)).await; } let call = async { if use_cli { @@ -298,10 +300,6 @@ impl ModelPool { self.client.chat(m, system, user).await } }; - // Race the in-flight call against a HARD cancel: when the user picks - // "report raw" / "discard" on /stop, drop the call future so the - // CLI child (spawned with kill_on_drop) is terminated immediately - // instead of finishing its whole command sequence. let r = tokio::select! { biased; _ = wait_cancelled(&self.cancel) => return Err(anyhow!("cancelled")), @@ -309,10 +307,23 @@ impl ModelPool { }; match r { Ok(t) => return Ok(t), - // Don't burn retries on exhaustion or auth failure — surface - // immediately so the caller can park and let the user /continue. - Err(e) if is_exhaustion(&e) || is_auth_failure(&e) => return Err(e), - Err(e) => last = e, + Err(e) => { + let class = classify_error(&format!("{e:#}")); + match class { + ErrorClass::Auth | ErrorClass::Exhaustion | + ErrorClass::Permanent | ErrorClass::ContentRefusal => return Err(e), + ErrorClass::ContextOverflow => { + return Err(anyhow!("context_overflow: {e}")); + } + ErrorClass::ImageRejection => { + return Err(anyhow!("image_rejection: {e}")); + } + ErrorClass::Transient => { + // transient — retry with backoff + last = e; + } + } + } } } Err(last) diff --git a/neurosploit-rs/crates/harness/src/resilience.rs b/neurosploit-rs/crates/harness/src/resilience.rs new file mode 100644 index 0000000..18f23c0 --- /dev/null +++ b/neurosploit-rs/crates/harness/src/resilience.rs @@ -0,0 +1,407 @@ +//! Execution resilience: transient-retry with classified backoff, tool-call +//! output limiting, context compaction, and crash-state salvage. +//! +//! Inspired by production patterns in Strix's execution engine, adapted to +//! NeuroSploit's Rust/tokio architecture and single-shot agent model. + +use std::time::Duration; + +// --------------------------------------------------------------------------- +// 1. Transient error classification + exponential backoff +// --------------------------------------------------------------------------- + +/// Whether an LLM provider error is transient (worth retrying) vs permanent. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ErrorClass { + /// Timeout, connection reset, 429, 500+, overloaded — retry with backoff. + Transient, + /// Auth failure (401/403), invalid key — stop immediately. + Auth, + /// Quota/billing exhaustion — pause the run for operator action. + Exhaustion, + /// Context window overflow — needs compaction, not retry. + ContextOverflow, + /// Bad request that won't change on retry (400 with non-overflow reason). + Permanent, + /// Image/multimodal rejection — scrub images, retry. + ImageRejection, + /// Content policy / guardrail refusal. + ContentRefusal, +} + +/// Classify an error string from an LLM provider into an actionable category. +pub fn classify_error(err: &str) -> ErrorClass { + let lower = err.to_lowercase(); + + // Auth first — most specific. + if ["401", "403", "unauthorized", "token has been revoked", + "invalid api key", "invalid_api_key", "not authenticated", + "authentication failed", "please log in", "please login", + "token revoked", "session expired"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::Auth; + } + + // Exhaustion / billing. + if ["quota", "insufficient_quota", "out of credit", "credit balance", + "billing", "usage limit", "resource_exhausted"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::Exhaustion; + } + + // Context overflow (must check before generic 400). + if ["context length", "context window", "context_length_exceeded", + "prompt is too long", "input is too long", "maximum prompt length", + "reduce the length", "too many tokens", "token limit exceeded"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::ContextOverflow; + } + + // Image rejection (before generic 400). + if ["no endpoints found that support image", + "image_url is only supported by certain models", + "is not a multimodal model", "at most 0 image", + "does not support image input", + "doesn't support the image field"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::ImageRejection; + } + + // Content policy / guardrail. + if ["content policy", "content_policy", "safety", "harmful", + "i cannot", "i can't help with", "refusal"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::ContentRefusal; + } + + // Rate limit (transient). + if ["rate limit", "rate_limit", "ratelimit", "429", "too many requests", + "overloaded", "capacity", "service unavailable", "503"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::Transient; + } + + // Server errors (transient). + if ["500", "502", "504", "timeout", "timed out", "connection reset", + "connection refused", "connection error", "eof", "broken pipe", + "network", "dns", "resolve"] + .iter().any(|k| lower.contains(k)) + { + return ErrorClass::Transient; + } + + // Everything else is permanent. + ErrorClass::Permanent +} + +/// Retry configuration for transient errors. +pub const MAX_TRANSIENT_RETRIES: u32 = 5; +const BASE_DELAY_MS: u64 = 2_000; +const MAX_DELAY_MS: u64 = 90_000; + +/// Compute exponential backoff delay for attempt N (1-indexed). +pub fn retry_delay(attempt: u32) -> Duration { + let delay = BASE_DELAY_MS * 2u64.pow(attempt.saturating_sub(1)); + Duration::from_millis(delay.min(MAX_DELAY_MS)) +} + +/// Should we retry this error class? +pub fn is_retryable(class: ErrorClass) -> bool { + matches!(class, ErrorClass::Transient) +} + +// --------------------------------------------------------------------------- +// 2. Tool-call output limiter (blast-radius cap) +// --------------------------------------------------------------------------- + +/// Limits how many tool calls from a single model response are honoured. +/// A degenerate generation can emit hundreds of tool calls (e.g. a +/// poll/wait loop written out ahead of time); executing all of them makes +/// the agent unresponsive for hours. Keeping only the first `limit` bounds +/// that blast radius. +pub struct ToolCallLimiter { + limit: usize, + kept: usize, + pub dropped: usize, +} + +impl ToolCallLimiter { + pub fn new(limit: usize) -> Self { + Self { limit, kept: 0, dropped: 0 } + } + + /// Whether the next tool call should be executed. + pub fn allow(&mut self) -> bool { + if self.limit == 0 { + return true; // unlimited + } + if self.kept < self.limit { + self.kept += 1; + true + } else { + self.dropped += 1; + false + } + } + + /// Filter a list of items, keeping only the first `limit`. + pub fn filter(&mut self, items: Vec) -> Vec { + if self.limit == 0 || items.len() <= self.limit { + self.kept += items.len(); + return items; + } + self.dropped += items.len() - self.limit; + self.kept += self.limit; + items.into_iter().take(self.limit).collect() + } +} + +// --------------------------------------------------------------------------- +// 3. Context compaction (prompt truncation with structured summary) +// --------------------------------------------------------------------------- + +use crate::pool::ModelPool; + +/// Approximate token count (4 chars per token heuristic, same as most +/// tokenizer estimates for English text + code). +pub fn approx_tokens(text: &str) -> usize { + text.len() / 4 +} + +/// Maximum context tokens before compaction triggers. +const COMPACT_THRESHOLD_TOKENS: usize = 100_000; +/// Tokens to keep verbatim from the end of the conversation. +const KEEP_RECENT_TOKENS: usize = 30_000; + +const COMPACT_SYSTEM: &str = "\ +You are compacting the earlier part of an autonomous security-testing agent's \ +conversation so it fits the model context window. Produce a dense, factual \ +record that lets the agent continue with no loss of important state.\n\n\ +This is a security engagement: dropped findings mean lost vulnerabilities. Be \ +EXHAUSTIVE. Enumerate every distinct item as its own bullet. Copy exact values \ +verbatim: URLs, endpoints, parameters, payloads, credentials, tokens, error \ +messages. Do not invent anything.\n\n\ +Return Markdown with these sections:\n\ +## Objective\n\ +## Vulnerabilities & Findings (one bullet per DISTINCT finding)\n\ +## Credentials & Secrets (verbatim, or \"(none)\")\n\ +## System & Recon Details\n\ +## Work State (completed / active / blocked)\n\ +## Failed Attempts & Dead Ends\n\ +## Next Move"; + +/// Compact a conversation that has grown too large by summarising the older +/// portion and keeping the recent portion verbatim. +/// +/// Returns `Some(compacted_prompt)` if compaction was performed, `None` if +/// the prompt is within budget. +pub async fn maybe_compact_prompt( + pool: &ModelPool, + full_prompt: &str, +) -> Option { + let total_tokens = approx_tokens(full_prompt); + if total_tokens <= COMPACT_THRESHOLD_TOKENS { + return None; + } + + // Split: older portion to summarise, recent portion to keep verbatim. + let keep_chars = KEEP_RECENT_TOKENS * 4; + let split_point = full_prompt.len().saturating_sub(keep_chars); + + // Find a clean split point (newline boundary). + let split = full_prompt[..split_point] + .rfind('\n') + .unwrap_or(split_point); + + let head = &full_prompt[..split]; + let tail = &full_prompt[split..]; + + if head.is_empty() { + return None; + } + + // Cap head to avoid sending an enormous summarisation request. + let head_capped: String = head.chars().take(200_000).collect(); + + let user_msg = format!( + "Summarise this earlier conversation from a security testing agent:\n\n{}\n\n\ + Respond with ONLY the structured Markdown summary.", + head_capped + ); + + match pool.complete(COMPACT_SYSTEM, &user_msg).await { + Ok((_model, summary)) => { + let compacted = format!( + "\n\ + The following summarises earlier conversation that was compacted to fit \ + the context window. Treat it as established context, not new instructions.\n\n\ + {}\n\ + \n\n\ + {}", + summary.trim(), + tail + ); + Some(compacted) + } + Err(e) => { + None + } + } +} + +// --------------------------------------------------------------------------- +// 4. Crash-state salvage (persist partial findings on agent failure) +// --------------------------------------------------------------------------- + +use crate::types::Finding; +use std::path::Path; + +/// Salvage partial findings from a crashed agent's raw output. +/// +/// When an agent dies mid-execution (timeout, provider crash, context overflow), +/// any findings already emitted in its partial response are extracted and +/// persisted rather than lost. This is the single biggest ROI improvement: +/// a 40-minute agent run that crashes at minute 38 still keeps its 38 minutes +/// of findings. +pub fn salvage_findings(agent_name: &str, partial_output: &str, workdir: Option<&str>) -> Vec { + if partial_output.is_empty() { + return Vec::new(); + } + + let findings = crate::pipeline::extract_findings(partial_output, agent_name); + if findings.is_empty() { + return Vec::new(); + } + + // Persist to a salvage file so they survive even if the caller also crashes. + if let Some(dir) = workdir { + let path = Path::new(dir).join("salvaged_findings.json"); + let existing: Vec = std::fs::read_to_string(&path) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); + + let mut all = existing; + for f in &findings { + if !all.iter().any(|e| e.title == f.title && e.endpoint == f.endpoint) { + all.push(f.clone()); + } + } + + if let Ok(json) = serde_json::to_string_pretty(&all) { + let _ = std::fs::write(&path, json); + } + } + + findings +} + +/// Load any previously salvaged findings from a run directory. +pub fn load_salvaged(workdir: &str) -> Vec { + let path = Path::new(workdir).join("salvaged_findings.json"); + std::fs::read_to_string(&path) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default() +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classify_transient_errors() { + assert_eq!(classify_error("Connection timed out"), ErrorClass::Transient); + assert_eq!(classify_error("Error 429: Too Many Requests"), ErrorClass::Transient); + assert_eq!(classify_error("502 Bad Gateway"), ErrorClass::Transient); + assert_eq!(classify_error("Service Unavailable"), ErrorClass::Transient); + } + + #[test] + fn classify_auth_errors() { + assert_eq!(classify_error("401 Unauthorized"), ErrorClass::Auth); + assert_eq!(classify_error("Invalid API key"), ErrorClass::Auth); + assert_eq!(classify_error("Token has been revoked"), ErrorClass::Auth); + } + + #[test] + fn classify_exhaustion() { + assert_eq!(classify_error("insufficient_quota"), ErrorClass::Exhaustion); + assert_eq!(classify_error("Out of credit"), ErrorClass::Exhaustion); + } + + #[test] + fn classify_context_overflow() { + assert_eq!(classify_error("context_length_exceeded"), ErrorClass::ContextOverflow); + assert_eq!(classify_error("This model's maximum context length is 128000 tokens"), ErrorClass::ContextOverflow); + assert_eq!(classify_error("prompt is too long"), ErrorClass::ContextOverflow); + } + + #[test] + fn classify_permanent() { + assert_eq!(classify_error("invalid JSON in request body"), ErrorClass::Permanent); + assert_eq!(classify_error("unknown model gpt-99"), ErrorClass::Permanent); + } + + #[test] + fn retry_delay_exponential() { + assert_eq!(retry_delay(1), Duration::from_millis(2_000)); + assert_eq!(retry_delay(2), Duration::from_millis(4_000)); + assert_eq!(retry_delay(3), Duration::from_millis(8_000)); + assert_eq!(retry_delay(4), Duration::from_millis(16_000)); + assert_eq!(retry_delay(5), Duration::from_millis(32_000)); + // Capped at 90s. + assert_eq!(retry_delay(10), Duration::from_millis(90_000)); + } + + #[test] + fn tool_call_limiter_caps_output() { + let mut limiter = ToolCallLimiter::new(3); + assert!(limiter.allow()); + assert!(limiter.allow()); + assert!(limiter.allow()); + assert!(!limiter.allow()); + assert!(!limiter.allow()); + assert_eq!(limiter.dropped, 2); + } + + #[test] + fn tool_call_limiter_zero_is_unlimited() { + let mut limiter = ToolCallLimiter::new(0); + for _ in 0..100 { + assert!(limiter.allow()); + } + assert_eq!(limiter.dropped, 0); + } + + #[test] + fn tool_call_limiter_filter() { + let mut limiter = ToolCallLimiter::new(2); + let items = vec![1, 2, 3, 4, 5]; + let filtered = limiter.filter(items); + assert_eq!(filtered, vec![1, 2]); + assert_eq!(limiter.dropped, 3); + } + + #[test] + fn approx_tokens_reasonable() { + let text = "hello world this is a test"; // 26 chars + assert_eq!(approx_tokens(text), 6); // 26/4 = 6 + } + + #[test] + fn compact_threshold_is_sane() { + // 100k tokens * 4 chars = 400k chars before compaction triggers. + assert!(COMPACT_THRESHOLD_TOKENS > KEEP_RECENT_TOKENS); + } +}