@vllm-sr/dsh 0.4.1 → 0.5.0

package/src/shield.js

@@ -1,11 +1,12 @@
 /**
  * The injection shield: screen tool output passage by passage, then zoom into
- * flagged passages to find the lines that carry the injection, and remove only
- * those. Lines the harness itself writes are masked while scoring and never
- * removed.
+ * flagged passages to find what carries the injection, and remove only that:
+ * the exact spans when a span model answers, else the lines. Lines the harness
+ * itself writes are masked while scoring and never removed.
  */
 import { ITEMS_PER_REQUEST, yes } from './client.js'
-import { SHIELD, SHIELD_FOCUS } from './questions.js'
+import { INJECTION_SPANS, SHIELD, SHIELD_FOCUS } from './questions.js'
+import { excerpts, placeSpans, replaceSpans, toSentences, unnumber } from './spans.js'
 
 /** Exact templates of the annotations DeepSeek Harness adds to shell output. */
 const HARNESS_LINES = [
@@ -50,6 +51,23 @@ async function scoreAll(client, model, questions, states, signal, limits = {}) {
   }
 }
 
+export const SPAN_MARK = '[System One removed injected text]'
+
+// The exact spans that give orders, in each flagged passage, from a span model; a passage it fails on gets none.
+async function injectedSpans({ client, model, about, passages, signal, timeoutMs, retries }) {
+  const found = passages.map(() => [])
+  let answeredBy
+  for (let start = 0; start < passages.length; start += ITEMS_PER_REQUEST) {
+    const batch = passages.slice(start, start + ITEMS_PER_REQUEST)
+    const response = await client.ask({ model, questions: INJECTION_SPANS, inputs: batch.map(passage => ({ state: { request: about, passage } })), signal, timeoutMs, retries })
+    answeredBy = response.model
+    batch.forEach((passage, offset) => {
+      found[start + offset] = placeSpans(response.spans[offset]?.injected, passage)
+    })
+  }
+  return { model: answeredBy ?? model, found }
+}
+
 /**
  * @param {object} options
  * @param {{ask: Function}} options.client
@@ -59,11 +77,18 @@ async function scoreAll(client, model, questions, states, signal, limits = {}) {
  * @param {number} options.threshold
  * @param {number} options.chunkChars
  * @param {AbortSignal} [options.signal]
- * @returns {Promise<{model: string, passages: number, maxP: number, flagged: object[], removed: Array<{start: number, end: number, p: number}>, removedLines: number, text: string}>}
+ * @param {{model: string, about: string}} [options.spans] - a span model that finds the injected text itself
+ * @returns {Promise<{model: string, passages: number, maxP: number, flagged: object[], removed: object[], removedLines: number,
+ *   removedSpans: number, removedChars: number, spanModel?: string, spanMs?: number, text: string}>}
  *   `maxP` is the highest injection probability of any passage
  */
-export async function screen({ client, model, prefix, text, threshold, chunkChars, signal, timeoutMs, retries }) {
+export async function screen({ client, model, prefix, text, threshold, chunkChars, signal, timeoutMs, retries, spans }) {
   const lines = text.split('\n')
+  const lineStart = []
+  for (let index = 0, at = 0; index < lines.length; index++) {
+    lineStart.push(at)
+    at += lines[index].length + 1
+  }
   const masked = lines.map(line => (isHarnessLine(line) ? '[harness note]' : line))
   const ranges = passageRanges(masked, chunkChars)
   const quote = (from, to, marks) => {
@@ -73,12 +98,48 @@ export async function screen({ client, model, prefix, text, threshold, chunkChar
     return `${prefix}\n"""\n${body.join('\n')}\n"""`
   }
   const first = await scoreAll(client, model, SHIELD, ranges.map(([from, to]) => quote(from, to)), signal, { timeoutMs, retries })
-  const flagged = ranges.map(([from, to], index) => ({ from, to, p: first.scores[index] })).filter(entry => entry.p >= threshold)
+  const flagged = ranges.map(([from, to], index) => ({ from, to, p: first.scores[index], text: lines.slice(from, to).join('\n') })).filter(entry => entry.p >= threshold)
+
+  // With a span model: the text that gives orders, in each flagged passage, grown to whole sentences. The passage is
+  // scored again without them, and a cut that leaves an injection behind, or a failure, leaves the lines below.
+  let spanModel
+  let spanMs
+  if (spans && flagged.length) {
+    const started = performance.now()
+    try {
+      // the span model reads a file's text without the read tool's line numbers; each cut is mapped back onto them
+      const plain = flagged.map(entry => unnumber(entry.text))
+      const answer = await injectedSpans({ client, model: spans.model, about: spans.about, passages: plain.map(view => view.text), signal, timeoutMs, retries })
+      spanModel = answer.model
+      const cuts = flagged.map((entry, index) => {
+        const grown = toSentences(plain[index].text, answer.found[index])
+        // the harness's own lines are never removed
+        if (!grown.length || grown.some(span => span.text.split('\n').some(isHarnessLine))) return undefined
+        entry.plain = { text: plain[index].text, spans: grown }
+        return grown.map(span => {
+          const start = plain[index].at(span.start)
+          const end = plain[index].at(span.end)
+          return { ...span, start, end, text: entry.text.slice(start, end) }
+        })
+      })
+      const tried = flagged.map((entry, index) => ({ entry, cut: cuts[index] })).filter(item => item.cut)
+      if (tried.length) {
+        const unmask = value => value.split('\n').map(line => (isHarnessLine(line) ? '[harness note]' : line)).join('\n')
+        const again = await scoreAll(client, model, SHIELD, tried.map(({ entry, cut }) => `${prefix}\n"""\n${unmask(replaceSpans(entry.text, cut, () => SPAN_MARK))}\n"""`), signal, { timeoutMs, retries })
+        tried.forEach(({ entry, cut }, index) => {
+          if (again.scores[index] < threshold) entry.spans = cut
+          else entry.spansLeft = again.scores[index]
+        })
+      }
+    } catch (error) {
+      if (signal?.aborted) throw error
+    }
+    spanMs = performance.now() - started
+  }
 
-  // Second look, only inside flagged passages: which groups of lines carry it?
+  // Otherwise a second look inside each flagged passage: which groups of lines carry it?
   const GROUP = 3
-  const removed = []
-  await Promise.all(flagged.map(async entry => {
+  await Promise.all(flagged.filter(entry => !entry.spans).map(async entry => {
     const groups = []
     for (let start = entry.from; start < entry.to; start += GROUP) groups.push([start, Math.min(entry.to, start + GROUP)])
     let hits = []
@@ -87,38 +148,47 @@ export async function screen({ client, model, prefix, text, threshold, chunkChar
       hits = zoom ? groups.map(([start, end], index) => ({ start, end, p: zoom.scores[index] })).filter(group => group.p >= threshold) : []
     }
     entry.lines = hits.length ? hits : [{ start: entry.from, end: entry.to, p: entry.p }]
-    removed.push(...entry.lines)
   }))
 
-  removed.sort((a, b) => a.start - b.start)
-  const merged = []
-  for (const range of removed) {
-    const last = merged.at(-1)
-    if (last && last.end >= range.start) {
-      last.end = Math.max(last.end, range.end)
-      last.p = Math.max(last.p, range.p)
-    } else {
-      merged.push({ ...range })
+  // Every removal as a range of the original text and what replaces it.
+  const lineEnd = index => (index < lines.length ? lineStart[index] : text.length)
+  const cuts = []
+  for (const entry of flagged) {
+    if (entry.spans) {
+      for (const span of entry.spans) {
+        const at = lineStart[entry.from]
+        cuts.push({ start: at + span.start, end: at + span.end, p: span.probability, replacement: SPAN_MARK, spans: 1, chars: span.end - span.start })
+      }
+      continue
+    }
+    const merged = []
+    for (const range of [...entry.lines].sort((a, b) => a.start - b.start)) {
+      const last = merged.at(-1)
+      if (last && last.end >= range.start) {
+        last.end = Math.max(last.end, range.end)
+        last.p = Math.max(last.p, range.p)
+      } else {
+        merged.push({ ...range })
+      }
+    }
+    for (const range of merged) {
+      const kept = lines.slice(range.start, range.end).filter(isHarnessLine)
+      const dropped = range.end - range.start - kept.length
+      if (dropped === 0) continue
+      const body = [`[System One removed ${dropped} line(s) here: they give orders to AI agents (p ${range.p.toFixed(2)})]`, ...kept].join('\n')
+      cuts.push({ start: lineStart[range.start], end: lineEnd(range.end), p: range.p, replacement: range.end < lines.length ? `${body}\n` : body, lines: dropped, startLine: range.start, endLine: range.end })
     }
   }
-
-  const out = []
-  let removedLines = 0
+  cuts.sort((a, b) => a.start - b.start)
+  let out = ''
   let cursor = 0
-  for (const range of merged) {
-    out.push(...lines.slice(cursor, range.start))
-    const kept = []
-    let dropped = 0
-    for (let index = range.start; index < range.end; index++) {
-      if (isHarnessLine(lines[index])) kept.push(lines[index])
-      else dropped++
-    }
-    if (dropped > 0) out.push(`[System One removed ${dropped} line(s) here: they give orders to AI agents (p ${range.p.toFixed(2)})]`)
-    out.push(...kept)
-    removedLines += dropped
-    cursor = range.end
+  for (const cut of cuts) {
+    out += text.slice(cursor, cut.start) + cut.replacement
+    cursor = cut.end
   }
-  out.push(...lines.slice(cursor))
+  out += text.slice(cursor)
+  const removedLines = cuts.reduce((sum, cut) => sum + (cut.lines ?? 0), 0)
+  const removedSpans = cuts.reduce((sum, cut) => sum + (cut.spans ?? 0), 0)
 
   return {
     model: first.model,
@@ -126,11 +196,17 @@ export async function screen({ client, model, prefix, text, threshold, chunkChar
     maxP: first.scores.reduce((max, p) => Math.max(max, p), 0),
     flagged: flagged.map(entry => ({
       p: entry.p,
-      text: lines.slice(entry.from, entry.to).join('\n'),
-      lines: entry.lines.map(range => ({ p: range.p, text: lines.slice(range.start, range.end).join('\n') })),
+      text: entry.text,
+      lines: (entry.lines ?? []).map(range => ({ p: range.p, text: lines.slice(range.start, range.end).join('\n') })),
+      spans: (entry.spans ? entry.plain.spans : []).map(span => ({ label: span.label, p: span.probability, text: span.text })),
+      excerpts: entry.spans ? excerpts(entry.plain.text, entry.plain.spans, { around: 60, max: 2 }) : [],
     })),
-    removed: merged,
+    removed: cuts.map(cut => (cut.lines ? { start: cut.startLine, end: cut.endLine, p: cut.p } : { span: true, start: cut.start, end: cut.end, p: cut.p })),
     removedLines,
-    text: removedLines > 0 ? out.join('\n') : text,
+    removedSpans,
+    removedChars: cuts.reduce((sum, cut) => sum + (cut.chars ?? 0), 0),
+    spanModel,
+    spanMs,
+    text: cuts.length > 0 ? out : text,
   }
 }

package/src/spans.js

@@ -0,0 +1,216 @@
+/**
+ * Span answers (Vela 3.0 models, contract 3.2): `spans.<id> = [{label, start, end, text, probability}]`, with offsets in
+ * Unicode code points into the text the question read, `end` exclusive. JavaScript strings index UTF-16 units, so every
+ * span is moved onto them and checked against its own text before it is used.
+ */
+
+/** UTF-16 index of each code point boundary of `text` (one more entry than code points). */
+function boundaries(text) {
+  const at = [0]
+  let index = 0
+  for (const char of text) {
+    index += char.length
+    at.push(index)
+  }
+  return at
+}
+
+/**
+ * The spans of one answer that really cut `text` where they say, as UTF-16 ranges, sorted and without overlaps (the
+ * more probable of two overlapping spans wins). A span whose offsets miss is found again by its text, nearest first.
+ * @param {Array<{label: string, start: number, end: number, text: string, probability: number}>} spans
+ * @param {string} text - what the question read
+ * @param {object} [options]
+ * @param {number} [options.min] - the lowest probability kept
+ * @returns {Array<{label: string, start: number, end: number, text: string, probability: number}>}
+ */
+export function placeSpans(spans, text, { min = 0 } = {}) {
+  if (!Array.isArray(spans) || typeof text !== 'string' || text.length === 0) return []
+  const at = boundaries(text)
+  const placed = []
+  for (const span of spans) {
+    if (!span || typeof span.text !== 'string' || span.text.length === 0) continue
+    const probability = typeof span.probability === 'number' && Number.isFinite(span.probability) ? span.probability : 0
+    if (probability < min) continue
+    let start = Number.isSafeInteger(span.start) && span.start >= 0 && span.start < at.length ? at[span.start] : -1
+    let end = Number.isSafeInteger(span.end) && span.end > 0 && span.end < at.length ? at[span.end] : -1
+    if (start < 0 || end <= start || text.slice(start, end) !== span.text) {
+      const near = start >= 0 ? start : 0
+      const after = text.indexOf(span.text, near)
+      const before = text.lastIndexOf(span.text, near)
+      const found = after < 0 ? before : before < 0 ? after : near - before <= after - near ? before : after
+      if (found < 0) continue
+      start = found
+      end = found + span.text.length
+    }
+    placed.push({ label: String(span.label ?? ''), start, end, text: span.text, probability })
+  }
+  placed.sort((a, b) => a.start - b.start || b.probability - a.probability)
+  const kept = []
+  for (const span of placed) {
+    const last = kept.at(-1)
+    if (last && span.start < last.end) {
+      if (span.probability > last.probability) kept[kept.length - 1] = span
+      continue
+    }
+    kept.push(span)
+  }
+  return kept
+}
+
+/** `text` with each span replaced by what `mask(span)` returns. */
+export function replaceSpans(text, spans, mask) {
+  let out = ''
+  let cursor = 0
+  for (const span of spans) {
+    out += text.slice(cursor, span.start) + mask(span)
+    cursor = span.end
+  }
+  return out + text.slice(cursor)
+}
+
+/** Personal data masked by its label, as `[EMAIL_ADDRESS]`. */
+export const maskLabel = span => `[${span.label || 'PII'}]`
+
+/**
+ * Up to `max` short excerpts of `text` around the spans, for Chat: plain parts and marked parts. With `hide`, a marked
+ * part carries only its label, so personal data never leaves the text it was found in.
+ * @returns {Array<Array<{text?: string, mark?: {label: string, probability: number, text?: string}}>>}
+ */
+export function excerpts(text, spans, { around = 48, max = 3, hide = false, clip = 160 } = {}) {
+  const groups = []
+  for (const span of spans) {
+    const from = Math.max(0, span.start - around)
+    const to = Math.min(text.length, span.end + around)
+    const last = groups.at(-1)
+    if (last && from <= last.to) {
+      last.to = Math.max(last.to, to)
+      last.spans.push(span)
+    } else {
+      groups.push({ from, to, spans: [span] })
+    }
+  }
+  const flat = value => value.replace(/\s+/g, ' ')
+  return groups.slice(0, max).map(group => {
+    const parts = []
+    let cursor = group.from
+    if (group.from > 0) parts.push({ text: '…' })
+    for (const span of group.spans) {
+      if (span.start > cursor) parts.push({ text: flat(text.slice(cursor, span.start)) })
+      const shown = flat(span.text)
+      parts.push({ mark: { label: span.label, probability: Math.round(span.probability * 1000) / 1000, ...(hide ? {} : { text: shown.length > clip ? `${shown.slice(0, clip - 1)}…` : shown }) } })
+      cursor = span.end
+    }
+    if (group.to > cursor) parts.push({ text: flat(text.slice(cursor, group.to)) })
+    if (group.to < text.length) parts.push({ text: '…' })
+    return parts
+  })
+}
+
+// merge ranges that touch or overlap, keeping the higher probability
+function merged(spans, text) {
+  const out = []
+  for (const span of [...spans].sort((a, b) => a.start - b.start)) {
+    const last = out.at(-1)
+    if (last && span.start <= last.end) {
+      last.end = Math.max(last.end, span.end)
+      last.probability = Math.max(last.probability, span.probability)
+      last.text = text.slice(last.start, last.end)
+    } else {
+      out.push({ ...span, text: text.slice(span.start, span.end) })
+    }
+  }
+  return out
+}
+
+/**
+ * Each span grown to the whole sentence it sits in: an injection is a sentence of orders, and a span model may mark
+ * only its loudest words. A sentence ends after `.`, `!` or `?`, at a blank line, and at a line that holds only a
+ * delimiter (`<!--`, `-->`, a code fence).
+ */
+export function toSentences(text, spans) {
+  const stops = [0, text.length]
+  for (const match of text.matchAll(/[.!?](?=\s)/g)) stops.push(match.index + 1)
+  for (const pattern of [/\n[ \t]*\n/g, /^[ \t]*(?:<!--|-->|```|~~~)[^\n]*$/gm]) {
+    for (const match of text.matchAll(pattern)) stops.push(match.index, match.index + match[0].length)
+  }
+  stops.sort((a, b) => a - b)
+  const grown = spans.map(span => {
+    let start = stops.filter(stop => stop <= span.start).pop()
+    const end = stops.find(stop => stop >= span.end)
+    while (start < span.start && /\s/.test(text[start])) start++
+    return { ...span, start, end }
+  })
+  return merged(grown, text)
+}
+
+/** Each span grown to the shell command it sits in (between `&&`, `||`, `;`, `|` and line breaks). */
+export function toCommands(command, spans) {
+  const stops = [0, command.length]
+  for (const match of command.matchAll(/&&|\|\||[;|\n]/g)) stops.push(match.index, match.index + match[0].length)
+  stops.sort((a, b) => a - b)
+  const grown = spans.map(span => {
+    let start = stops.filter(stop => stop <= span.start).pop()
+    let end = stops.find(stop => stop >= span.end)
+    while (start < end && /\s/.test(command[start])) start++
+    while (end > start && /\s/.test(command[end - 1])) end--
+    return { ...span, start, end }
+  })
+  return merged(grown, command)
+}
+
+/**
+ * A text without the line numbers some tools print before each line (`12: `), and the way back: `at(index)` is where a
+ * position of the plain text sits in the numbered one. A text whose lines are not numbered comes back as it is.
+ */
+export function unnumber(text) {
+  const lines = text.split('\n')
+  const filled = lines.filter(line => line.length > 0)
+  if (filled.length === 0 || filled.filter(line => /^\d+: ?/.test(line)).length < filled.length * 0.8) return { text, at: index => index }
+  const plainStarts = []
+  const numberedStarts = []
+  const cut = []
+  let plain = 0
+  let numbered = 0
+  const out = lines.map(line => {
+    const bare = line.replace(/^\d+: ?/, '')
+    plainStarts.push(plain)
+    numberedStarts.push(numbered)
+    cut.push(line.length - bare.length)
+    plain += bare.length + 1
+    numbered += line.length + 1
+    return bare
+  })
+  const at = index => {
+    let line = 0
+    while (line + 1 < plainStarts.length && plainStarts[line + 1] <= index) line++
+    return numberedStarts[line] + cut[line] + (index - plainStarts[line])
+  }
+  return { text: out.join('\n'), at }
+}
+
+/** How many spans of each label, most first. */
+export function labelCounts(spans) {
+  const counts = {}
+  for (const span of spans) counts[span.label || 'PII'] = (counts[span.label || 'PII'] ?? 0) + 1
+  return Object.fromEntries(Object.entries(counts).sort((a, b) => b[1] - a[1]))
+}
+
+/**
+ * Cut `text` into windows of at most `size` UTF-16 units that end at line breaks where they can, for questions whose
+ * target is limited. @returns {Array<{start: number, text: string}>}
+ */
+export function windows(text, size) {
+  const out = []
+  let start = 0
+  while (start < text.length) {
+    let end = Math.min(text.length, start + size)
+    if (end < text.length) {
+      const nl = text.lastIndexOf('\n', end)
+      if (nl > start + size / 2) end = nl + 1
+    }
+    out.push({ start, text: text.slice(start, end) })
+    start = end
+  }
+  return out
+}