Expose trace progress and per-record truncation for bounded highlight recovery #68

diff --git a/src/engine.rs b/src/engine.rs
index d805b83..bf444c6 100644
--- a/src/engine.rs
+++ b/src/engine.rs
@@ -4685,4 +4685,304 @@ mod tests {
 		};
 		assert_eq!(error.code, "E_INCOMPLETE_CREATE");
 	}
+	fn highlight_issue_fixture() -> (Engine, SearchRequest) {
+		let mut cfg = config();
+		cfg.identity.surface_terms = true;
+		let engine = Engine::open(RamDirectory::create(), &cfg).unwrap();
+		let request = SearchRequest {
+			expression: expression("needle", SearchMode::Any, vec!["description".to_owned()]),
+			candidate_ids: None,
+			filter: None,
+			offset: 0,
+			limit: 10,
+			exact_total: true,
+			budget_milliseconds: 30_000,
+			traversal: false,
+			after: None,
+		};
+		(engine, request)
+	}
+	fn highlight_issue_record(id: &str, n: usize) -> TraceRecord {
+		TraceRecord {
+			id: id.to_owned(),
+			fields: vec![("description".to_owned(), vec!["needle ".repeat(n)])],
+		}
+	}
+	fn highlight_issue_counts(result: &TraceResult) -> Vec<(String, usize)> {
+		result
+			.records
+			.iter()
+			.map(|r| (r.id.clone(), r.values.iter().map(|v| v.spans.len()).sum()))
+			.collect()
+	}
+	#[test]
+	fn highlight_issue_baseline_batch_and_singleton() {
+		let (engine, request) = highlight_issue_fixture();
+		let records = [
+			highlight_issue_record("a", 600),
+			highlight_issue_record("b", 600),
+			highlight_issue_record("c", 600),
+		];
+		let batch = engine.trace_matches(&request, &records, None).unwrap();
+		assert!(!batch.complete);
+		assert_eq!(
+			highlight_issue_counts(&batch),
+			vec![("a".to_owned(), 600), ("b".to_owned(), 424)]
+		);
+		for row in &records {
+			let result = engine.trace_matches(&request, std::slice::from_ref(row), None).unwrap();
+			assert!(result.complete);
+			assert_eq!(highlight_issue_counts(&result)[0].1, 600);
+		}
+		let singleton = engine
+			.trace_matches(&request, &[highlight_issue_record("large", 1100)], None)
+			.unwrap();
+		assert!(!singleton.complete);
+		assert_eq!(highlight_issue_counts(&singleton)[0].1, 1024);
+		println!("BASELINE batch=a:600,b:424,c:omitted; individual=600 complete; oversized=1024 incomplete");
+	}
+	#[test]
+	fn highlight_issue_prototype_progress_preserves_completed_records() {
+		let (engine, request) = highlight_issue_fixture();
+		let records = [
+			highlight_issue_record("a", 600),
+			highlight_issue_record("b", 600),
+			highlight_issue_record("c", 600),
+		];
+		let mut offset = 0;
+		let mut seen = Vec::new();
+		let mut calls = 0;
+		while offset < records.len() {
+			let result = engine.probe_trace_progress(&request, &records[offset..], None).unwrap();
+			assert!(result.next_record > 0);
+			assert!(result.states.iter().all(|(_, spans, analysis)| !spans && !analysis));
+			seen.extend(highlight_issue_counts(&result.trace));
+			offset += result.next_record;
+			calls += 1;
+		}
+		assert_eq!(
+			seen,
+			vec![("a".to_owned(), 600), ("b".to_owned(), 600), ("c".to_owned(), 600)]
+		);
+		assert_eq!(calls, 3);
+		println!(
+			"PROTOTYPE progress: 3 calls, 1800 spans, completed records emitted once; not a one-extra-call guarantee"
+		);
+	}
+	#[test]
+	fn highlight_issue_prototype_singleton_truncation_is_explicit() {
+		let (engine, request) = highlight_issue_fixture();
+		let result = engine
+			.probe_trace_progress(&request, &[highlight_issue_record("large", 1100)], None)
+			.unwrap();
+		assert_eq!(result.next_record, 1);
+		assert!(!result.trace.complete);
+		assert_eq!(result.states, vec![("large".to_owned(), true, false)]);
+		assert_eq!(highlight_issue_counts(&result.trace)[0].1, 1024);
+		println!("PROTOTYPE singleton: bounded1024 spans; explicit span truncation; source analysis not truncated");
+	}
+	#[test]
+	fn highlight_issue_prototype_no_match_and_analysis_unknown_are_distinct() {
+		let (engine, request) = highlight_issue_fixture();
+		let unmatched = TraceRecord {
+			id: "none".to_owned(),
+			fields: vec![("description".to_owned(), vec!["stone trail".to_owned()])],
+		};
+		let result = engine
+			.probe_trace_progress(&request, &[unmatched, highlight_issue_record("a", 1)], None)
+			.unwrap();
+		assert_eq!(result.next_record, 2);
+		assert_eq!(result.states.len(), 2);
+		assert!(result.trace.complete);
+		assert_eq!(highlight_issue_counts(&result.trace), vec![("a".to_owned(), 1)]);
+		let unknown = TraceRecord {
+			id: "unknown".to_owned(),
+			fields: vec![(
+				"description".to_owned(),
+				vec![format!("{}needle", "x ".repeat(MAX_TRACE_TOKENS_PER_VALUE))],
+			)],
+		};
+		let result = engine.probe_trace_progress(&request, &[unknown], None).unwrap();
+		assert!(!result.trace.complete);
+		assert!(result.trace.records.is_empty());
+		assert_eq!(result.states, vec![("unknown".to_owned(), false, true)]);
+	}
+	#[test]
+	fn highlight_issue_prototype_preserves_deadline() {
+		let (engine, request) = highlight_issue_fixture();
+		let deadline = Some(Instant::now() - std::time::Duration::from_millis(1));
+		let error = match engine.probe_trace_progress(&request, &[highlight_issue_record("a", 1)], deadline) {
+			Ok(_) => panic!("expired deadline accepted"),
+			Err(error) => error,
+		};
+		assert_eq!(error.code, "E_TIMEOUT");
+	}
+}
+
+#[cfg(test)]
+struct ProbeTrace {
+	trace: TraceResult,
+	next_record: usize,
+	states: Vec<(String, bool, bool)>,
+}
+#[cfg(test)]
+impl Engine {
+	fn probe_trace_progress(
+		&self,
+		request: &SearchRequest,
+		records: &[TraceRecord],
+		deadline: Option<Instant>,
+	) -> Result<ProbeTrace> {
+		check_deadline(deadline)?;
+		if let Some(candidate_ids) = &request.candidate_ids {
+			for id in candidate_ids {
+				validate_record_id(id)?;
+			}
+		}
+		for record in records {
+			validate_record_id(&record.id)?;
+		}
+		let clause = match &request.expression {
+			SearchExpression::Clause(clause) => clause,
+			_ => {
+				return Err(FulltextError::invalid(
+					"traceMatches does not support boolean expressions",
+				))
+			}
+		};
+		let selected = self.selected_fields(&clause.fields)?;
+		self.require_surface_fields(&selected)?;
+		let selected_names = selected.iter().map(|field| field.name.as_str()).collect::<HashSet<_>>();
+		let candidates = request
+			.candidate_ids
+			.as_ref()
+			.map(|ids| ids.iter().map(String::as_str).collect::<HashSet<_>>());
+		if candidates.as_ref().is_some_and(HashSet::is_empty) {
+			return Ok(ProbeTrace {
+				trace: TraceResult {
+					complete: true,
+					records: Vec::new(),
+				},
+				next_record: records.len(),
+				states: Vec::new(),
+			});
+		}
+		let plan = self.trace_plan(clause)?;
+		let mut complete = true;
+		let mut remaining_spans = MAX_TRACE_SPANS;
+		let mut response_bytes = 3usize;
+		let mut matched_records = Vec::new();
+		let mut next_record = 0usize;
+		let mut states = Vec::new();
+		let mut analyzed_source_analyzer = self.index_analyzer.clone();
+		let mut surface_source_analyzer = self.surface_index_analyzer.clone();
+		for (record_index, record) in records.iter().enumerate() {
+			check_deadline(deadline)?;
+			if remaining_spans == 0 {
+				complete = false;
+				break;
+			}
+			if candidates
+				.as_ref()
+				.is_some_and(|candidates| !candidates.contains(record.id.as_str()))
+			{
+				next_record = record_index + 1;
+				continue;
+			}
+			let mut seen_fields = HashSet::new();
+			let mut pending_values = Vec::new();
+			let mut found = HashSet::new();
+			let mut record_span_budget = remaining_spans;
+			let mut record_truncated = false;
+			let mut record_analysis_truncated = false;
+			for (field, source_values) in &record.fields {
+				if !seen_fields.insert(field.as_str()) {
+					return Err(FulltextError::invalid(format!("duplicate trace field {field}")));
+				}
+				if !self.field_lookup.contains_key(field) {
+					return Err(FulltextError::invalid(format!("unknown trace field {field}")));
+				}
+				if !selected_names.contains(field.as_str()) {
+					continue;
+				}
+				for (value_index, value) in source_values.iter().enumerate() {
+					check_deadline(deadline)?;
+					let (spans, terms, spans_truncated, analysis_truncated) = self.trace_value(
+						&plan,
+						value,
+						record_span_budget,
+						deadline,
+						&mut analyzed_source_analyzer,
+						&mut surface_source_analyzer,
+					)?;
+					record_truncated |= spans_truncated || analysis_truncated;
+					record_analysis_truncated |= analysis_truncated;
+					found.extend(terms);
+					if spans.is_empty() {
+						continue;
+					}
+					record_span_budget = record_span_budget.saturating_sub(spans.len());
+					pending_values.push((field.as_str(), value_index as u32, spans));
+				}
+			}
+			if record_truncated && !record_analysis_truncated && next_record > 0 {
+				complete = false;
+				break;
+			}
+			next_record = record_index + 1;
+			states.push((
+				record.id.clone(),
+				record_truncated && !record_analysis_truncated,
+				record_analysis_truncated,
+			));
+			complete &= !record_analysis_truncated;
+			if plan.record_matches(&found) {
+				complete &= !record_truncated;
+				if pending_values.is_empty() {
+					continue;
+				}
+				let record_header_bytes = 6usize.saturating_add(record.id.len());
+				let mut record_bytes = 0usize;
+				let mut values = Vec::new();
+				for (field, value_index, spans) in pending_values {
+					let value_header_bytes = 10usize.saturating_add(field.len());
+					let available = MAX_SEARCH_RESPONSE_BYTES
+						.saturating_sub(response_bytes)
+						.saturating_sub(record_header_bytes)
+						.saturating_sub(record_bytes);
+					let byte_limited_spans = available.saturating_sub(value_header_bytes) / 8;
+					let keep = spans.len().min(remaining_spans).min(byte_limited_spans);
+					if keep < spans.len() {
+						complete = false;
+					}
+					if keep == 0 {
+						continue;
+					}
+					values.push(TraceValueMatch {
+						field: field.to_owned(),
+						value_index,
+						spans: spans.into_iter().take(keep).collect(),
+					});
+					remaining_spans -= keep;
+					record_bytes = record_bytes.saturating_add(value_header_bytes + keep * 8);
+				}
+				if values.is_empty() {
+					continue;
+				}
+				response_bytes = response_bytes.saturating_add(record_header_bytes + record_bytes);
+				matched_records.push(TraceRecordMatch {
+					id: record.id.clone(),
+					values,
+				});
+			}
+		}
+		Ok(ProbeTrace {
+			trace: TraceResult {
+				complete,
+				records: matched_records,
+			},
+			next_record,
+			states,
+		})
+	}
 }