CalvinBackup/NeuroSploit
1
0 Fork 0
mirror of https://github.com/CyberSecurityUP/NeuroSploit.git synced 2026-10-04 15:06:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: robust LLM response handling & JSON extraction (#46)

Browse Source 
* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Corentin GoetghebeurandClaude Opus 4.8 authored and GitHub committed 2026-10-02 18:48:50 -03:00
1 parent 5d4e7e0347
commit 01eac8f0c3
7 files changed
+845 -129
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

neurosploit-rs/Cargo.lock
Generated
+60
View File
@@ -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"
neurosploit-rs/crates/harness/Cargo.toml
+1
View File
@@ -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
neurosploit-rs/crates/harness/src/json_extract.rs
+308
View File
@@ -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<Value> {
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<Value> {
let s = slice.trim();
if s.is_empty() {
return None;
}
if let Ok(v) = serde_json::from_str::<Value>(s) {
return Some(v);
}
if let Ok(v) = json5::from_str::<Value>(s) {
return Some(v);
}
if let Some(repaired) = close_truncated(s) {
if let Ok(v) = serde_json::from_str::<Value>(&repaired) {
return Some(v);
}
if let Ok(v) = json5::from_str::<Value>(&repaired) {
return Some(v);
}
}
None
}
/// Candidate JSON slices, in the order we should try them.
fn candidates(text: &str) -> Vec<String> {
let mut out: Vec<String> = 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<String>, Option<usize>) {
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<usize> {
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<String> {
let bytes = slice.as_bytes();
let mut stack: Vec<u8> = 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<u8>)> = 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());
}
}
neurosploit-rs/crates/harness/src/lib.rs
+1
View File
@@ -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;
neurosploit-rs/crates/harness/src/models.rs
+236 -9
View File
@@ -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::<u64>().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::<u32>().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<String> {
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::<serde_json::Value>(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::<Vec<_>>()
.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::<String>())
}
}
#[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());
}
}
neurosploit-rs/crates/harness/src/pipeline.rs
+235 -102
View File
@@ -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<Finding>, Vec<String>) {
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::<serde_json::Value>(&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<Strin
}
}
/// Keys a model wraps a names/strings array under when it returns an object
/// instead of the bare array we asked for — `{"agents":[…]}`, `{"selected":[…]}`.
/// Same deviation as [`FINDINGS_WRAPPER_KEYS`], seen on the agent-selection reply;
/// without unwrapping, a wrapped selection reads as empty and the run silently
/// falls back to RL ranking.
const STRING_ARRAY_WRAPPER_KEYS: &[&str] =
&["agents", "selected", "selection", "names", "agent_names", "chosen", "list", "items"];
/// Fields that carry the string inside an array element that is an *object*
/// (`[{"name":"sqli"}]`) rather than a bare string (`["sqli"]`).
const STRING_ELEMENT_KEYS: &[&str] = &["name", "agent", "id", "value"];
/// Pull a list of strings out of a model reply, tolerating the shapes models use
/// when they stray from "reply with a JSON array of strings": a wrapper object
/// around the array, and array elements that are objects carrying the string in a
/// `name`/`agent`/`id` field.
fn parse_string_array(text: &str) -> Vec<String> {
match (text.find('['), text.rfind(']')) {
(Some(a), Some(b)) if b > a => serde_json::from_str::<Vec<String>>(&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<Finding>)]) -> 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::<serde_json::Value>(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<serde_json::Value> {
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<Finding> {
let mut candidates: Vec<String> = 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::<serde_json::Value>(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<serde_json::Value> = 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)]
neurosploit-rs/crates/harness/src/prosecutor.rs
+4 -18
View File
@@ -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<ProsecutorVerdict> {
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::<ProsecutorVerdict>(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::<ProsecutorVerdict>(v).ok()
}
/// Fold the prosecutor's reading into the finding's claim set.