From 01eac8f0c33b57b49815ef9d1939a6fd01f8e1bb Mon Sep 17 00:00:00 2001 From: Corentin Goetghebeur <72756186+CorentinGoet@users.noreply.github.com> Date: Fri, 2 Oct 2026 23:48:50 +0200 Subject: [PATCH] fix: robust LLM response handling & JSON extraction (#46) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(pipeline): robust LLM JSON extraction (json5 + truncation repair) Model replies that did not exactly match the expected JSON syntax were either dropped silently or surfaced as "[extract_findings] ... JSON parse failed" / "... no JSON array/object found". Both came from the same two weak stages in extract_findings: a greedy first-'['-to-last-']' span that captured prose, and a salvage pass that only stripped trailing commas. Add a shared, string/escape-aware extractor (crates/harness/src/json_extract.rs): - locate balanced [..]/{..} regions, ignoring brackets inside prose/strings, preferring fenced blocks (last wins); - parse leniently: serde_json first, then json5 (trailing commas, comments, single quotes, unquoted keys); - repair token-limit truncation by closing the open structure, keeping the complete findings instead of discarding the whole batch. Route extract_findings, reported_nothing, extract_chain, parse_string_array and prosecutor::parse_verdict through it. Make the diagnostic tail() char-boundary-safe (the old slice could panic on UTF-8). Add regression tests for single quotes/comments, capitalised ```JSON fences, truncated arrays and pure prose. Co-Authored-By: Claude Opus 4.8 * fix(models): robust LLM response handling + higher token/timeout limits Harden the OpenAI-compatible chat client against the empty-content and parse failures hit with reasoning models (GLM/DeepSeek via OpenRouter) during whitebox runs: - Accept message `content` as a string, an array of content parts, or a `reasoning_content` fallback; surface `finish_reason` and empty-choices errors instead of an opaque "no content in response". - Stop masking mid-stream body-read failures as a bogus "EOF while parsing"; report read timeouts and empty bodies explicitly, and reassemble SSE-framed responses some gateways return unrequested. - Raise reasoning-model max_tokens to 32768 and the HTTP timeout to 300s; both overridable via NEUROSPLOIT_MAX_TOKENS / NEUROSPLOIT_HTTP_TIMEOUT. Adds unit tests for content extraction, token sizing, and SSE reassembly. Cargo.lock syncs the json5 entry from the prior extraction commit. Co-Authored-By: Claude Opus 4.8 * fix(pipeline): unwrap findings/selection replies wrapped in an object The json5 + truncation work made JSON *parsing* robust, but the *shape* handling after it still dropped data when a model wrapped its answer in an object instead of returning the bare array we asked for. Most visible on black-box runs, where a long tool-use turn ends with the model narrating into a report object. extract_findings treated any object as ONE finding, so a real batch returned as `{"findings":[…]}` (or `{"vulnerabilities":[…]}`, …) became a single title-less "finding", was filtered out, and surfaced to the operator as "returned text but 0 parseable findings" while the findings were lost. Add findings_items() to normalise the shape: an array is the list; an object with a title is one bare finding; otherwise an object wrapping a known findings key unwraps to that array. reported_nothing() now recognises the same wrapper keys so an empty `{"vulnerabilities":[]}` reads as an honest negative. Two more consumers of the same class: - parse_string_array (agent selection) accepted only a bare array of strings, so a wrapped `{"agents":[…]}` or elements-as-objects `[{"name":"sqli"}]` silently fell back to RL ranking. Now unwraps the wrapper and pulls the string from object elements. - extract_chain hard-coded the "findings" key for its object branch, dropping the sibling `loot` under any other wrapper key. Now checks all wrapper keys. Add regression tests for each shape. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 --- neurosploit-rs/Cargo.lock | 60 ++++ neurosploit-rs/crates/harness/Cargo.toml | 1 + .../crates/harness/src/json_extract.rs | 308 ++++++++++++++++ neurosploit-rs/crates/harness/src/lib.rs | 1 + neurosploit-rs/crates/harness/src/models.rs | 245 ++++++++++++- neurosploit-rs/crates/harness/src/pipeline.rs | 337 ++++++++++++------ .../crates/harness/src/prosecutor.rs | 22 +- 7 files changed, 845 insertions(+), 129 deletions(-) create mode 100644 neurosploit-rs/crates/harness/src/json_extract.rs diff --git a/neurosploit-rs/Cargo.lock b/neurosploit-rs/Cargo.lock index a3f6100..ee16cc6 100644 --- a/neurosploit-rs/Cargo.lock +++ b/neurosploit-rs/Cargo.lock @@ -855,6 +855,17 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "json5" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b0db21af676c1ce64250b5f40f3ce2cf27e4e47cb91ed91eb6fe9350b430c1" +dependencies = [ + "pest", + "pest_derive", + "serde", +] + [[package]] name = "libc" version = "0.2.186" @@ -952,6 +963,7 @@ dependencies = [ "base64", "futures", "hmac", + "json5", "regex", "reqwest", "serde", @@ -1029,6 +1041,48 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pest" +version = "2.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45d3aca230fad2e6f6317ca0a72724338c4960cb97168a85cdee66df4a9a21a8" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "284b60557f2c4a2e72ad3f2d34d42685a2fa4a6a61d0d2a10c0ae2a5e916c2cf" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d9d1f08a115309ee99268cf85e5228e0e56aa9caf8841ec12866b6be07c3109" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pest_meta" +version = "2.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed93ba1a9ffcca32130a5188701c81c0c49cf00d4b7c5007d5148951d743adcb" +dependencies = [ + "pest", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1798,6 +1852,12 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/neurosploit-rs/crates/harness/Cargo.toml b/neurosploit-rs/crates/harness/Cargo.toml index 1525018..ea1b8c1 100644 --- a/neurosploit-rs/crates/harness/Cargo.toml +++ b/neurosploit-rs/crates/harness/Cargo.toml @@ -11,6 +11,7 @@ path = "src/lib.rs" [dependencies] serde.workspace = true serde_json.workspace = true +json5 = "0.4" tokio.workspace = true reqwest.workspace = true anyhow.workspace = true diff --git a/neurosploit-rs/crates/harness/src/json_extract.rs b/neurosploit-rs/crates/harness/src/json_extract.rs new file mode 100644 index 0000000..c8bced1 --- /dev/null +++ b/neurosploit-rs/crates/harness/src/json_extract.rs @@ -0,0 +1,308 @@ +//! Robust extraction of JSON from language-model replies. +//! +//! Models rarely return exactly the JSON we asked for. They wrap it in prose, +//! fence it (sometimes as ```JSON with a capitalised tag), add trailing commas +//! or `// comments`, quote with single quotes, and — most damaging — get cut +//! off mid-array by a token limit. The strict [`serde_json`] path throws on +//! every one of these and the agent's findings are lost. The two failures seen +//! on live runs were the two ends of this: `JSON parse failed` (a region was +//! found but was not valid JSON) and `no JSON array/object found` (the greedy +//! `[`…`]` span never located one). +//! +//! This module recovers both in three stages: +//! +//! 1. **Locate** every *balanced* `[...]` / `{...}` region in the reply, scanning +//! with string- and escape-awareness so brackets inside prose (`[low] …`) or +//! inside a string value never fool the matcher. Fenced blocks are preferred, +//! and later regions beat earlier ones (models narrate, then answer). +//! 2. **Parse leniently**: strict `serde_json` first (the fast path for good +//! input), then [`json5`] for the common relaxations — trailing commas, +//! comments, single quotes and unquoted keys. +//! 3. **Repair truncation**: when a region never closes, drop the dangling +//! partial element and append the closers the open structure still needs, so +//! a reply cut off after three complete findings still yields three. +//! +//! Every candidate is parse-checked, so a region that cannot be made into valid +//! JSON is simply skipped — the function never returns a half-parsed value, and +//! `None` means the reply genuinely held no recoverable JSON. + +use serde_json::Value; + +/// Parse a model reply into a JSON value, trying each located candidate in +/// best-first order and returning the first that parses. `None` means the reply +/// contained no recoverable JSON (genuine prose, a refusal, or an empty reply). +pub fn parse_reply(text: &str) -> Option { + for cand in candidates(text) { + if let Some(v) = parse_lenient(&cand) { + return Some(v); + } + } + None +} + +/// Parse one slice: strict JSON, then json5, then a truncation-repair retry. +pub fn parse_lenient(slice: &str) -> Option { + let s = slice.trim(); + if s.is_empty() { + return None; + } + if let Ok(v) = serde_json::from_str::(s) { + return Some(v); + } + if let Ok(v) = json5::from_str::(s) { + return Some(v); + } + if let Some(repaired) = close_truncated(s) { + if let Ok(v) = serde_json::from_str::(&repaired) { + return Some(v); + } + if let Ok(v) = json5::from_str::(&repaired) { + return Some(v); + } + } + None +} + +/// Candidate JSON slices, in the order we should try them. +fn candidates(text: &str) -> Vec { + let mut out: Vec = Vec::new(); + // Fenced blocks first, last fence winning — models draft, then give the + // final answer in the last block. + for block in fenced_blocks(text).into_iter().rev() { + let (regions, _) = regions_and_truncation(block); + if regions.is_empty() { + out.push(block.to_string()); + } else { + for r in regions.into_iter().rev() { + out.push(r); + } + } + } + // Then balanced regions anywhere in the reply, again last-wins. + let (regions, trunc) = regions_and_truncation(text); + for r in regions.into_iter().rev() { + out.push(r); + } + // If the scan hit a structure that never closed, hand the whole truncated + // tail (from where it opened) to `parse_lenient`, which will try to close it. + // Doing this for the *outer* opener is what keeps a cut-off array of findings + // whole, instead of salvaging only its first complete element. + if let Some(start) = trunc { + out.push(text[start..].to_string()); + } else if out.is_empty() { + if let Some(start) = text.find(|c| c == '[' || c == '{') { + out.push(text[start..].to_string()); + } + } + out.dedup(); + out +} + +/// The body of every ```fenced``` block, in order. A language tag on the opening +/// line (```json, ```JSON, ```json5) is irrelevant: the balanced scanner finds +/// the JSON inside regardless of it — which is exactly what the old +/// `starts_with('[')` check got wrong for a capitalised tag. +fn fenced_blocks(text: &str) -> Vec<&str> { + let mut out = Vec::new(); + let mut rest = text; + while let Some(open) = rest.find("```") { + let after = &rest[open + 3..]; + let Some(close) = after.find("```") else { break }; + out.push(after[..close].trim()); + rest = &after[close + 3..]; + } + out +} + +/// Every balanced top-level `[...]` / `{...}` region, in order of appearance, +/// plus — if the scan reaches an opener that never closes — the byte index where +/// that truncated structure begins. Brackets inside string literals (and escaped +/// quotes) are ignored, so prose like `[low] …` or a payload string containing +/// `]` does not break matching. +/// +/// When an opener does not close, scanning stops there: everything after it is +/// *inside* that unclosed structure, not an independent region, so the inner +/// objects of a cut-off array must not be mistaken for the whole answer. +fn regions_and_truncation(text: &str) -> (Vec, Option) { + let bytes = text.as_bytes(); + let mut out = Vec::new(); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'[' || bytes[i] == b'{' { + match scan_balanced(bytes, i) { + Some(end) => { + out.push(text[i..end].to_string()); + i = end; + continue; + } + None => return (out, Some(i)), + } + } + i += 1; + } + (out, None) +} + +/// From an opener at `start`, return the byte index just past its matching +/// closer, or `None` if the structure never closes (a truncated reply). JSON +/// structural bytes are all ASCII and never appear inside a UTF-8 multibyte +/// sequence, so byte scanning is UTF-8-safe and the returned index lands on a +/// char boundary. +fn scan_balanced(bytes: &[u8], start: usize) -> Option { + let mut depth: i32 = 0; + let mut in_str = false; + let mut esc = false; + let mut i = start; + while i < bytes.len() { + let c = bytes[i]; + if in_str { + if esc { + esc = false; + } else if c == b'\\' { + esc = true; + } else if c == b'"' { + in_str = false; + } + } else { + match c { + b'"' => in_str = true, + b'[' | b'{' => depth += 1, + b']' | b'}' => { + depth -= 1; + if depth == 0 { + return Some(i + 1); + } + } + _ => {} + } + } + i += 1; + } + None +} + +/// Repair a reply cut off mid-structure: cut back to the last complete element +/// (a closed container, or the position just before the last comma) and append +/// the closers the still-open containers need. Returns `None` when nothing is +/// open to close, or there is no complete element to cut back to. +fn close_truncated(slice: &str) -> Option { + let bytes = slice.as_bytes(); + let mut stack: Vec = Vec::new(); + let mut in_str = false; + let mut esc = false; + // (byte index to cut to, open-container stack at that point). + let mut cut: Option<(usize, Vec)> = None; + let mut i = 0; + while i < bytes.len() { + let c = bytes[i]; + if in_str { + if esc { + esc = false; + } else if c == b'\\' { + esc = true; + } else if c == b'"' { + in_str = false; + } + } else { + match c { + b'"' => in_str = true, + b'[' => stack.push(b']'), + b'{' => stack.push(b'}'), + b']' | b'}' => { + stack.pop(); + // A complete element just closed: cutting here keeps it. + cut = Some((i + 1, stack.clone())); + } + // Between elements: cut before the comma, dropping whatever + // partial element follows it. + b',' => cut = Some((i, stack.clone())), + _ => {} + } + } + i += 1; + } + let (idx, open) = cut?; + if open.is_empty() { + return None; + } + let mut repaired = slice[..idx].trim_end().to_string(); + if repaired.ends_with(',') { + repaired.pop(); + } + for close in open.iter().rev() { + repaired.push(*close as char); + } + Some(repaired) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn strict_json_is_untouched() { + let v = parse_reply(r#"[{"title":"x"}]"#).unwrap(); + assert!(v.is_array()); + } + + #[test] + fn prose_brackets_before_the_answer_are_ignored() { + // `[low] …` is a balanced `[...]` region but not valid JSON; the real + // answer is the fenced block. + let t = "[low] cookie missing Secure\n\n```json\n[{\"title\":\"real\"}]\n```"; + let v = parse_reply(t).unwrap(); + assert_eq!(v[0]["title"], "real"); + } + + #[test] + fn capitalised_fence_tag_still_parses() { + let t = "```JSON\n[{\"title\":\"x\"}]\n```"; + assert!(parse_reply(t).unwrap().is_array()); + } + + #[test] + fn last_block_wins() { + let t = "```json\n[{\"title\":\"draft\"}]\n```\nthen\n```json\n[{\"title\":\"final\"}]\n```"; + assert_eq!(parse_reply(t).unwrap()[0]["title"], "final"); + } + + #[test] + fn trailing_comma_comments_and_single_quotes() { + let t = "```json5\n[ {'title': 'x'}, ] // note\n```"; + let v = parse_reply(t).unwrap(); + assert_eq!(v[0]["title"], "x"); + } + + #[test] + fn bracket_inside_a_string_value() { + let t = r#"[{"title":"he said ]"}]"#; + assert_eq!(parse_reply(t).unwrap()[0]["title"], "he said ]"); + } + + #[test] + fn truncated_array_keeps_the_complete_elements() { + // Cut off by a token limit mid-way through the third object. + let t = r#"[{"title":"a"},{"title":"b"},{"title":"#; + let v = parse_reply(t).unwrap(); + assert_eq!(v.as_array().unwrap().len(), 2); + assert_eq!(v[1]["title"], "b"); + } + + #[test] + fn truncated_unclosed_string_in_value() { + let t = r#"[{"title":"a"},{"title":"bbb"#; + let v = parse_reply(t).unwrap(); + assert_eq!(v.as_array().unwrap().len(), 1); + assert_eq!(v[0]["title"], "a"); + } + + #[test] + fn pure_prose_yields_none() { + assert!(parse_reply("I could not find any injectable parameters.").is_none()); + } + + #[test] + fn empty_array_is_recovered_not_lost() { + assert!(parse_reply("```json\n[]\n```").unwrap().as_array().unwrap().is_empty()); + } +} diff --git a/neurosploit-rs/crates/harness/src/lib.rs b/neurosploit-rs/crates/harness/src/lib.rs index aed84c6..781733c 100644 --- a/neurosploit-rs/crates/harness/src/lib.rs +++ b/neurosploit-rs/crates/harness/src/lib.rs @@ -25,6 +25,7 @@ pub mod inbox; pub mod integrations; pub mod integrity; pub mod internal; +pub mod json_extract; pub mod knowledge_graph; pub mod memory; pub mod policy; diff --git a/neurosploit-rs/crates/harness/src/models.rs b/neurosploit-rs/crates/harness/src/models.rs index 1a811f0..e73e141 100644 --- a/neurosploit-rs/crates/harness/src/models.rs +++ b/neurosploit-rs/crates/harness/src/models.rs @@ -127,10 +127,21 @@ pub struct ChatClient { http: reqwest::Client, } +/// HTTP request timeout (seconds) for an API chat completion. Large reasoning +/// responses (GLM/DeepSeek, raised `max_tokens`) can take minutes, so the +/// default is generous; override with `NEUROSPLOIT_HTTP_TIMEOUT` (seconds). +fn http_timeout_secs() -> u64 { + std::env::var("NEUROSPLOIT_HTTP_TIMEOUT") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|&v| v > 0) + .unwrap_or(300) +} + impl ChatClient { pub fn new() -> Self { let http = reqwest::Client::builder() - .timeout(Duration::from_secs(120)) + .timeout(Duration::from_secs(http_timeout_secs())) .build() .unwrap_or_else(|_| reqwest::Client::new()); ChatClient { http } @@ -170,7 +181,7 @@ impl ChatClient { }; let body = serde_json::json!({ "model": m.model, - "max_tokens": 4096, + "max_tokens": max_tokens_for(&m.model), "temperature": 0.2, "messages": [ {"role": "system", "content": system}, @@ -190,21 +201,71 @@ impl ChatClient { anyhow!("{} connection error: {}", p.key, e) } } else if e.is_timeout() { - anyhow!("{} request timed out (120s) for model '{}' — model may be too large for available memory", p.key, m.model) + anyhow!("{} request timed out ({}s) for model '{}' — raise NEUROSPLOIT_HTTP_TIMEOUT, or the model may be too large/slow for this prompt", p.key, http_timeout_secs(), m.model) } else { anyhow!("{} request error: {}", p.key, e) } })?; let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + // Read the body explicitly: a failure here (connection dropped mid-stream, + // or the request timeout firing while a large response is still streaming) + // must surface as itself, not get silently flattened to "" and then + // reappear downstream as a bogus "EOF while parsing" JSON error. + let text = resp.text().await.map_err(|e| { + if e.is_timeout() { + anyhow!("{} response body read timed out ({}s) for model '{}' — large/slow response was cut off mid-stream; raise NEUROSPLOIT_HTTP_TIMEOUT or lower max_tokens", p.key, http_timeout_secs(), m.model) + } else { + anyhow!("{} response body read failed (status {}) for model '{}': {}", p.key, status, m.model, e) + } + })?; if !status.is_success() { return Err(anyhow!("{} returned {}: {}", p.key, status, truncate(&text, 200))); } - let v: serde_json::Value = serde_json::from_str(&text)?; - let content = v["choices"][0]["message"]["content"] - .as_str() - .ok_or_else(|| anyhow!("no content in response"))?; - Ok(content.to_string()) + if text.trim().is_empty() { + return Err(anyhow!("{} returned an empty body (status {}) for model '{}'", p.key, status, m.model)); + } + let v: serde_json::Value = match serde_json::from_str(&text) { + Ok(v) => v, + Err(e) => { + // Some gateways emit Server-Sent Events even when stream wasn't + // requested. Reassemble the answer from the `data:` frames before + // giving up; only then report an unparseable body (with a snippet, + // so a truncated/HTML/error page is actually diagnosable). + if let Some(sse) = parse_sse_content(&text) { + if !sse.trim().is_empty() { + return Ok(sse); + } + } + return Err(anyhow!( + "{} returned unparseable body ({} bytes, model '{}'): {} — body starts: {}", + p.key, text.len(), m.model, e, truncate(text.trim_start(), 300) + )); + } + }; + let choice = &v["choices"][0]; + if choice.is_null() { + // Some providers signal a soft failure with an empty `choices` and an + // `error` object in an otherwise-200 body — surface that, not a blank. + let err = v["error"]["message"].as_str().unwrap_or("empty choices"); + return Err(anyhow!("{} returned no choices: {}", p.key, truncate(err, 200))); + } + let content = extract_message_text(&choice["message"]); + if content.trim().is_empty() { + // 2xx but no usable text. The usual cause is a reasoning model that + // spent the whole token budget on its reasoning stream and was cut + // off (finish_reason="length") before emitting the answer. + let finish = choice["finish_reason"].as_str().unwrap_or("unknown"); + let hint = if finish == "length" { + " — raised max_tokens may help, or the model is reasoning-heavy for this prompt size" + } else { + "" + }; + return Err(anyhow!( + "{} returned empty content (finish_reason={finish}, model='{}'){hint}", + p.key, m.model + )); + } + Ok(content) } } @@ -835,6 +896,116 @@ impl Default for ChatClient { } } +/// Default output budget for a chat completion, keyed off the model family. +/// +/// Reasoning models (GLM, DeepSeek, Qwen "thinking", the o-series, …) spend a +/// large, hidden chunk of their output budget on an internal reasoning stream +/// *before* the answer. With a 4K cap and a big prompt — e.g. whitebox code +/// review, which inlines a whole source bundle — they routinely hit +/// `finish_reason: "length"` with an empty `content`, surfacing as the +/// "empty content" error. Give those families more headroom; leave the rest at +/// the cheaper default. Override for all models with `NEUROSPLOIT_MAX_TOKENS`. +fn max_tokens_for(model: &str) -> u32 { + if let Some(v) = std::env::var("NEUROSPLOIT_MAX_TOKENS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|&v| v > 0) + { + return v; + } + let m = model.to_ascii_lowercase(); + // Strip a provider/route prefix like "anthropic/" or "z-ai/" so the family + // match works for OpenRouter/LiteLLM-style ids too. + let fam = m.rsplit('/').next().unwrap_or(&m); + let reasoning = fam.contains("glm") + || fam.contains("deepseek") + || fam.contains("qwen") + || fam.contains("kimi") + || fam.contains("reason") + || fam.contains("think") + || fam.contains("-r1") + || fam.contains("-r2") + // OpenAI o-series (o1/o3/o4…) and "sol" reasoning variants. + || fam.starts_with('o') && fam.chars().nth(1).is_some_and(|c| c.is_ascii_digit()) + || fam.contains("-sol"); + if reasoning { + 32768 + } else { + 4096 + } +} + +/// Reassemble assistant text from a Server-Sent Events body — the streaming +/// chat-completions format (`data: {json}\n\n` frames, ending with +/// `data: [DONE]`). Some gateways return this even when `stream` wasn't +/// requested, which makes a whole-body `serde_json::from_str` fail. Each frame's +/// delta lives at `choices[0].delta.content` (final frames may use `message` +/// instead). Returns `None` when the body isn't SSE at all. +fn parse_sse_content(body: &str) -> Option { + let mut saw_frame = false; + let mut out = String::new(); + for line in body.lines() { + let line = line.trim_start(); + let Some(payload) = line.strip_prefix("data:") else { continue }; + let payload = payload.trim(); + if payload.is_empty() || payload == "[DONE]" { + saw_frame = true; + continue; + } + let Ok(v) = serde_json::from_str::(payload) else { continue }; + saw_frame = true; + let choice = &v["choices"][0]; + if let Some(delta) = choice.get("delta") { + out.push_str(&extract_message_text(delta)); + } + // Non-streaming frame delivered over SSE (some gateways do this). + if choice.get("message").is_some() { + out.push_str(&extract_message_text(&choice["message"])); + } + } + if saw_frame { + Some(out) + } else { + None + } +} + +/// Pull the assistant text out of an OpenAI-compatible `message`, tolerating the +/// three response shapes seen across providers: +/// 1. `content` is a plain string (classic OpenAI). +/// 2. `content` is an array of parts `[{"type":"text","text":"…"}, …]` +/// (some proxies and Z.ai/GLM in certain modes). +/// 3. `content` is null but the text lives in `reasoning_content` +/// (reasoning models like GLM / DeepSeek, often when cut off early). +/// Returns an empty string when none carries usable text, so the caller can +/// report a precise, diagnosable error instead of a bare "no content". +fn extract_message_text(msg: &serde_json::Value) -> String { + if let Some(s) = msg["content"].as_str() { + if !s.trim().is_empty() { + return s.to_string(); + } + } + if let Some(arr) = msg["content"].as_array() { + let joined: String = arr + .iter() + .filter_map(|part| part["text"].as_str().or_else(|| part.as_str())) + .collect::>() + .join(""); + if !joined.trim().is_empty() { + return joined; + } + } + // Last resort: a reasoning model that never finalized an answer. The chain + // of thought often still contains the JSON we want, and downstream + // extraction is robust to surrounding prose. + if let Some(s) = msg["reasoning_content"].as_str() { + if !s.trim().is_empty() { + return s.to_string(); + } + } + String::new() +} + fn truncate(s: &str, n: usize) -> String { // Truncate by CHARACTERS, never bytes — slicing `&s[..n]` panics when `n` // lands inside a multi-byte char (e.g. '—'). That panic was crashing agent @@ -845,3 +1016,59 @@ fn truncate(s: &str, n: usize) -> String { format!("{}…", s.chars().take(n).collect::()) } } + +#[cfg(test)] +mod response_tests { + use super::*; + use serde_json::json; + + #[test] + fn reasoning_models_get_more_headroom() { + for m in ["glm-5.3", "glm-4.6", "deepseek-v4.1", "z-ai/glm-5.3", "qwen3.8-max", "kimi-k3", "o3", "gpt-6-sol"] { + assert_eq!(max_tokens_for(m), 32768, "{m} should be treated as reasoning-heavy"); + } + } + + #[test] + fn plain_models_keep_default() { + for m in ["gpt-5.5", "claude-opus-5-5", "anthropic/claude-opus-4-8", "gemini-3-pro", "grok-4.7", "meta-llama/llama-3.3-70b-instruct"] { + assert_eq!(max_tokens_for(m), 4096, "{m} should keep the default cap"); + } + } + + #[test] + fn extracts_plain_string_content() { + let msg = json!({"content": "hello"}); + assert_eq!(extract_message_text(&msg), "hello"); + } + + #[test] + fn extracts_array_content_parts() { + let msg = json!({"content": [{"type": "text", "text": "foo"}, {"type": "text", "text": "bar"}]}); + assert_eq!(extract_message_text(&msg), "foobar"); + } + + #[test] + fn falls_back_to_reasoning_content() { + let msg = json!({"content": serde_json::Value::Null, "reasoning_content": "the answer"}); + assert_eq!(extract_message_text(&msg), "the answer"); + } + + #[test] + fn reassembles_sse_stream() { + let body = "data: {\"choices\":[{\"delta\":{\"content\":\"foo\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"bar\"}}]}\n\ndata: [DONE]\n\n"; + assert_eq!(parse_sse_content(body).as_deref(), Some("foobar")); + } + + #[test] + fn non_sse_body_returns_none() { + assert!(parse_sse_content("{\"choices\":[]}").is_none()); + assert!(parse_sse_content("plain text").is_none()); + } + + #[test] + fn empty_everywhere_yields_empty() { + let msg = json!({"content": " "}); + assert!(extract_message_text(&msg).trim().is_empty()); + } +} diff --git a/neurosploit-rs/crates/harness/src/pipeline.rs b/neurosploit-rs/crates/harness/src/pipeline.rs index 6a44040..6b05148 100644 --- a/neurosploit-rs/crates/harness/src/pipeline.rs +++ b/neurosploit-rs/crates/harness/src/pipeline.rs @@ -1552,17 +1552,16 @@ async fn chain_from_seed(pool: &ModelPool, target: &str, directives: &str, recon /// Parse a chain agent reply into (new findings, loot). Accepts the object form /// `{"findings":[...],"loot":[...]}` and falls back to a bare findings array. fn extract_chain(text: &str, agent: &str) -> (Vec, Vec) { - if let (Some(a), Some(b)) = (text.find('{'), text.rfind('}')) { - if b > a { - if let Ok(serde_json::Value::Object(o)) = serde_json::from_str::(&text[a..=b]) { - if o.contains_key("findings") { - let findings = o.get("findings").map(|v| extract_findings(&v.to_string(), agent)).unwrap_or_default(); - let loot = o.get("loot").and_then(|v| v.as_array()) - .map(|arr| arr.iter().filter_map(|x| x.as_str().map(|s| s.to_string())).collect()) - .unwrap_or_default(); - return (findings, loot); - } - } + if let Some(serde_json::Value::Object(o)) = crate::json_extract::parse_reply(text) { + // A chain reply wrapping its findings array (under any of the keys models + // use, not just "findings") also carries `loot` as a sibling — pull both. + let findings_arr = FINDINGS_WRAPPER_KEYS.iter().find_map(|k| o.get(*k)); + if findings_arr.is_some() { + let findings = findings_arr.map(|v| extract_findings(&v.to_string(), agent)).unwrap_or_default(); + let loot = o.get("loot").and_then(|v| v.as_array()) + .map(|arr| arr.iter().filter_map(|x| x.as_str().map(|s| s.to_string())).collect()) + .unwrap_or_default(); + return (findings, loot); } } (extract_findings(text, agent), vec![]) @@ -1668,11 +1667,47 @@ async fn typesafe_prune_agents(recon: &str, catalog: &[Agent], chosen: Vec Vec { - match (text.find('['), text.rfind(']')) { - (Some(a), Some(b)) if b > a => serde_json::from_str::>(&text[a..=b]).unwrap_or_default(), - _ => vec![], - } + let arr = match crate::json_extract::parse_reply(text) { + Some(serde_json::Value::Array(a)) => a, + Some(serde_json::Value::Object(o)) => { + match STRING_ARRAY_WRAPPER_KEYS + .iter() + .find_map(|k| o.get(*k).and_then(|v| v.as_array())) + { + Some(a) => a.clone(), + None => return vec![], + } + } + _ => return vec![], + }; + arr.iter() + .filter_map(|v| match v { + serde_json::Value::String(s) => Some(s.trim().to_string()), + serde_json::Value::Object(o) => STRING_ELEMENT_KEYS + .iter() + .find_map(|k| o.get(*k).and_then(|x| x.as_str())) + .map(|s| s.trim().to_string()), + _ => None, + }) + .filter(|s| !s.is_empty()) + .collect() } /// Fallback agent selection when the LLM selector fails: score each agent by @@ -2858,107 +2893,99 @@ fn transcript_of(raw: &[(String, String, Vec)]) -> String { /// so we parse leniently into `Value` and coerce every field. /// Did the agent explicitly report an empty result? /// -/// Accepts a bare `[]`, a fenced ```json block containing one, and the common -/// `{"findings": []}` wrapper — all three mean "I looked and found nothing". +/// Accepts a bare `[]`, a fenced ```json block containing one, and an empty +/// findings wrapper (`{"findings":[]}`, `{"vulnerabilities":[]}`, … — see +/// [`FINDINGS_WRAPPER_KEYS`]) — all mean "I looked and found nothing". fn reported_nothing(text: &str) -> bool { - let mut t = text.trim(); - // Take the last fenced block when there is one; models narrate first and - // put the machine-readable answer at the end. - if let Some(start) = t.rfind("```") { - if let Some(open) = t[..start].rfind("```") { - let inner = &t[open + 3..start]; - let inner = inner.strip_prefix("json").unwrap_or(inner); - t = inner.trim(); - } - } - let t = t.trim_start_matches("```json").trim_start_matches("```").trim_end_matches("```").trim(); - if t == "[]" { + if text.trim() == "[]" { return true; } - serde_json::from_str::(t) - .map(|v| match &v { - serde_json::Value::Array(a) => a.is_empty(), - serde_json::Value::Object(o) => o.get("findings").and_then(|f| f.as_array()).map(|a| a.is_empty()).unwrap_or(false), - _ => false, - }) - .unwrap_or(false) + // Same lenient extractor the finding parser uses, so "nothing found" and + // "here are the findings" are decided from the exact same value — they can + // never disagree about which region of the reply is the answer. + match crate::json_extract::parse_reply(text) { + Some(serde_json::Value::Array(a)) => a.is_empty(), + Some(serde_json::Value::Object(o)) => FINDINGS_WRAPPER_KEYS + .iter() + .find_map(|k| o.get(*k).and_then(|f| f.as_array())) + .map(|a| a.is_empty()) + .unwrap_or(false), + _ => false, + } } -/// Every ```fenced``` block in the text, in order. -fn fenced_blocks(text: &str) -> Vec<&str> { - let mut out = Vec::new(); - let mut rest = text; - while let Some(open) = rest.find("```") { - let after = &rest[open + 3..]; - let Some(close) = after.find("```") else { break }; - let inner = &after[..close]; - let inner = inner.strip_prefix("json").unwrap_or(inner); - out.push(inner.trim()); - rest = &after[close + 3..]; +/// Last `n` characters of `s`, on a char boundary (diagnostics only). +fn tail(s: &str, n: usize) -> String { + let total = s.chars().count(); + s.chars().skip(total.saturating_sub(n)).collect() +} + +/// Keys a model wraps its findings array under when it ignores "reply with ONLY +/// a JSON array" and returns an object instead — `{"findings":[…]}`, +/// `{"vulnerabilities":[…]}`, etc. Seen most on black-box runs, where the reply +/// follows a long tool-use turn and the model narrates into a report object. None +/// of these names collide with a finding's own fields, so unwrapping is safe. +const FINDINGS_WRAPPER_KEYS: &[&str] = + &["findings", "vulnerabilities", "vulns", "results", "issues"]; + +/// Normalise a parsed reply into the list of candidate finding objects. +/// +/// - An array *is* the list. +/// - An object carrying a non-empty `title` is a single bare finding object. +/// - Otherwise an object wrapping one of [`FINDINGS_WRAPPER_KEYS`] unwraps to that +/// array — without this a real batch returned as `{"findings":[…]}` is treated +/// as one title-less "finding", dropped, and surfaces to the operator as +/// "returned text but 0 parseable findings" while the findings are lost. +/// - Any other object is passed through as a single finding object (so a bare +/// finding using a non-standard title key still reaches the field coercion). +fn findings_items(val: serde_json::Value) -> Vec { + match val { + serde_json::Value::Array(a) => a, + serde_json::Value::Object(o) => { + let has_title = o + .get("title") + .and_then(|v| v.as_str()) + .map(|t| !t.trim().is_empty()) + .unwrap_or(false); + if !has_title { + for k in FINDINGS_WRAPPER_KEYS { + if let Some(serde_json::Value::Array(a)) = o.get(*k) { + return a.clone(); + } + } + } + vec![serde_json::Value::Object(o)] + } + _ => vec![], } - out } /// Pull the findings array out of a model's reply. /// -/// The naive "first `[` to last `]`" span is wrong whenever the agent narrates -/// before answering: a real reply here opened with the prose line -/// `[low] Antiforgery cookie missing Secure flag`, so the span started inside -/// prose, failed to parse, and the agent's actual findings were thrown away. -/// Fenced blocks are tried first (last one wins — models narrate, then answer), -/// and the span is only a fallback. +/// Locating and parsing the JSON is delegated to [`crate::json_extract`], which +/// is string-aware (prose brackets like `[low] …` no longer start the span), +/// prefers the last fenced block, and tolerates the deviations that used to sink +/// a whole batch — fenced/capitalised tags, trailing commas, comments, single +/// quotes, and replies truncated by a token limit. The two failures seen on live +/// 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 { - let mut candidates: Vec = Vec::new(); - for b in fenced_blocks(text).into_iter().rev() { - if b.starts_with('[') || b.starts_with('{') { - candidates.push(b.to_string()); - } - } - if let (Some(a), Some(b)) = (text.find('['), text.rfind(']')) { - if b > a { - candidates.push(text[a..=b].to_string()); - } - } - if let (Some(a), Some(b)) = (text.find('{'), text.rfind('}')) { - if b > a { - candidates.push(text[a..=b].to_string()); - } - } - let slice: String = match candidates.iter().find(|c| serde_json::from_str::(c).is_ok()).cloned() { - Some(good) => good, - None => match candidates.into_iter().next() { - // Nothing parsed: keep the best guess so the salvage pass below - // still gets a shot at a trailing-comma mistake. - Some(first) => first, - None => { - if !text.trim().is_empty() && text.trim() != "[]" { - eprintln!("[extract_findings] agent {agent}: model returned text but no JSON array/object found (len={}); raw tail: {:?}", - text.len(), &text[text.len().saturating_sub(200)..]); - } - return vec![]; - } - }, - }; - let slice: &str = &slice; - let val: serde_json::Value = match serde_json::from_str(slice) { - Ok(v) => v, - Err(e) => { - eprintln!("[extract_findings] agent {agent}: JSON parse failed: {e}; slice head: {:?}", - &slice[..slice.len().min(300)]); - // Attempt to salvage: strip trailing comma before ] (common LLM mistake) - let fixed = slice.replace(",]", "]").replace(",}", "}"); - match serde_json::from_str(&fixed) { - Ok(v) => v, - Err(_) => return vec![], + let val = match crate::json_extract::parse_reply(text) { + Some(v) => v, + None => { + let t = text.trim(); + if !t.is_empty() && t != "[]" { + eprintln!( + "[extract_findings] agent {agent}: no parseable JSON in model reply (len={}); raw tail: {:?}", + text.len(), + tail(text, 200) + ); } + return vec![]; } }; - let items: Vec = match val { - serde_json::Value::Array(a) => a, - serde_json::Value::Object(_) => vec![val], - _ => return vec![], - }; - items + findings_items(val) .into_iter() .filter_map(|it| { let o = it.as_object()?; @@ -4017,6 +4044,112 @@ mod extraction_tests { let f = extract_findings("```json\n[{\"title\":\"X\",\"severity\":\"Low\"},]\n```", "a"); assert_eq!(f.len(), 1); } + + /// `JSON parse failed` on a live run: single-quoted keys/values and a `//` + /// comment — none of which strict serde accepts, all of which json5 does. + #[test] + fn single_quotes_and_comments_are_recovered() { + let f = extract_findings("```json\n[ {'title': 'Reflected XSS', 'severity': 'High'} ] // done\n```", "a"); + assert_eq!(f.len(), 1); + assert_eq!(f[0].title, "Reflected XSS"); + } + + /// A capitalised fence tag used to fail the old `starts_with('[')` gate and + /// surface as `no JSON array/object found`. + #[test] + fn a_capitalised_fence_tag_is_handled() { + let f = extract_findings("```JSON\n[{\"title\":\"Open redirect\",\"severity\":\"Medium\"}]\n```", "a"); + assert_eq!(f.len(), 1); + assert_eq!(f[0].title, "Open redirect"); + } + + /// Token-limit truncation mid-way through the third finding: the two complete + /// ones must survive instead of the whole batch being discarded. + #[test] + fn a_truncated_array_keeps_the_complete_findings() { + let text = "[{\"title\":\"A\",\"severity\":\"Low\"},{\"title\":\"B\",\"severity\":\"Low\"},{\"title\":\"C"; + let f = extract_findings(text, "a"); + assert_eq!(f.len(), 2); + assert_eq!(f[1].title, "B"); + } + + /// Pure prose (a refusal or narration with no JSON) is not a finding and not + /// a parse error — it yields nothing. + #[test] + fn pure_prose_yields_no_findings() { + assert!(extract_findings("I could not identify any injectable parameters.", "a").is_empty()); + } + + /// Black-box models often ignore "reply with ONLY a JSON array" and wrap the + /// batch in `{"findings":[…]}`. That object has no `title`, so the old + /// object-as-one-finding path dropped it and reported "0 parseable findings" + /// while real findings were lost. It must now unwrap to its array. + #[test] + fn a_findings_wrapper_object_is_unwrapped() { + let text = "```json\n{\"findings\":[{\"title\":\"IDOR on /orders\",\"severity\":\"High\"},{\"title\":\"Reflected XSS\",\"severity\":\"Medium\"}]}\n```"; + let f = extract_findings(text, "a"); + assert_eq!(f.len(), 2); + assert_eq!(f[0].title, "IDOR on /orders"); + } + + /// Alternate wrapper keys models use for the same shape. + #[test] + fn a_vulnerabilities_wrapper_object_is_unwrapped() { + let f = extract_findings("{\"vulnerabilities\":[{\"title\":\"SQLi\",\"severity\":\"Critical\"}]}", "a"); + assert_eq!(f.len(), 1); + assert_eq!(f[0].title, "SQLi"); + } + + /// A single bare finding object (with its own `title`) is still one finding — + /// unwrapping must not steal a nested array it happens to carry. + #[test] + fn a_bare_finding_object_stays_one_finding() { + let f = extract_findings("{\"title\":\"Open redirect\",\"severity\":\"Low\",\"repro_steps\":[\"curl …\"]}", "a"); + assert_eq!(f.len(), 1); + assert_eq!(f[0].title, "Open redirect"); + } + + /// An empty wrapper is an honest negative, not a malformed reply. + #[test] + fn an_empty_wrapper_is_reported_nothing_not_a_parse_failure() { + assert!(reported_nothing("{\"findings\":[]}")); + assert!(reported_nothing("```json\n{\"vulnerabilities\": []}\n```")); + assert!(extract_findings("{\"findings\":[]}", "a").is_empty()); + } + + /// The bare-array happy path for agent selection. + #[test] + fn a_string_array_parses_plain() { + assert_eq!(parse_string_array("[\"sqli\",\"xss\"]"), vec!["sqli", "xss"]); + } + + /// Same wrapper-object deviation as findings: `{"agents":[…]}` must not read as + /// an empty selection (which silently drops the model's choice to RL ranking). + #[test] + fn a_wrapped_string_array_is_unwrapped() { + assert_eq!(parse_string_array("{\"agents\":[\"sqli\",\"idor\"]}"), vec!["sqli", "idor"]); + assert_eq!(parse_string_array("```json\n{\"selected\": [\"ssrf\"]}\n```"), vec!["ssrf"]); + } + + /// Models sometimes answer with objects per element instead of bare strings. + #[test] + fn array_elements_that_are_objects_yield_their_name() { + assert_eq!( + parse_string_array("[{\"name\":\"sqli\",\"why\":\"x\"},{\"agent\":\"xss\"}]"), + vec!["sqli", "xss"] + ); + } + + /// A chain reply wrapping findings under a non-`findings` key still yields both + /// the findings and the sibling loot. + #[test] + fn extract_chain_unwraps_any_findings_key_with_loot() { + let text = "{\"vulnerabilities\":[{\"title\":\"RCE\",\"severity\":\"Critical\"}],\"loot\":[\"root creds\"]}"; + let (f, loot) = extract_chain(text, "chain"); + assert_eq!(f.len(), 1); + assert_eq!(f[0].title, "RCE"); + assert_eq!(loot, vec!["root creds"]); + } } #[cfg(test)] diff --git a/neurosploit-rs/crates/harness/src/prosecutor.rs b/neurosploit-rs/crates/harness/src/prosecutor.rs index e8caf8c..2ecd0f7 100644 --- a/neurosploit-rs/crates/harness/src/prosecutor.rs +++ b/neurosploit-rs/crates/harness/src/prosecutor.rs @@ -82,24 +82,10 @@ impl ProsecutorVerdict { /// Parse the prosecutor's reply, tolerating the fences and preamble models add. pub fn parse_verdict(text: &str) -> Option { - let mut candidates: Vec<&str> = Vec::new(); - // A fenced block is the machine-readable answer when there is one. - let mut rest = text; - let mut blocks: Vec<&str> = Vec::new(); - while let Some(open) = rest.find("```") { - let after = &rest[open + 3..]; - let Some(close) = after.find("```") else { break }; - let inner = after[..close].trim_start_matches("json").trim(); - blocks.push(inner); - rest = &after[close + 3..]; - } - candidates.extend(blocks.into_iter().rev()); - if let (Some(a), Some(b)) = (text.find('{'), text.rfind('}')) { - if b > a { - candidates.push(&text[a..=b]); - } - } - candidates.into_iter().find_map(|c| serde_json::from_str::(c).ok()) + // The shared extractor handles the fences, preamble and minor syntax drift + // (trailing commas, comments, single quotes) that models add to this object. + let v = crate::json_extract::parse_reply(text)?; + serde_json::from_value::(v).ok() } /// Fold the prosecutor's reading into the finding's claim set.