fix(retrieval): return only units with body text, filling open slots from the next-ranked evidence

internal/api/server.go

@@ -636,27 +636,14 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if body.MaxSections > 0 && len(ids) > body.MaxSections {
-		ids = ids[:body.MaxSections]
-	}
-
-	enriched := make([]sectionWithContent, 0, len(ids))
-	for _, id := range ids {
-		sec := t.FindByID(id)
-		if sec == nil {
-			continue
-		}
-		var content string
-		if sec.ContentRef != "" {
-			rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef)
-			if err == nil {
-				raw, _ := io.ReadAll(rc)
-				_ = rc.Close() // best-effort close
-				content = string(raw)
-			}
-		}
-		enriched = append(enriched, sectionWithContent{sec: sec, content: content})
+	// Every unit carries body text, and max_sections caps pages and
+	// sections together: a heading-only section, or one the evidence pages
+	// already carry, gives its slot to the next-ranked evidence (HAL-3160).
+	res := selResult
+	if res == nil {
+		res = &retrieval.Result{SelectedIDs: ids}
 	}
+	pages, enriched, runnerUpPages, unitStats := d.assembleUnits(r.Context(), t, res, body.MaxSections)
 
 	// Optional: content-aware re-rank pass. One LLM call that scores
 	// each loaded section against the query and re-orders the slice
@@ -677,18 +664,21 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
 		runSpansConcurrent(r.Context(), extractor, body.Query, enriched, d.AnswerSpan.MaxConcurrency, d.Logger)
 	}
 
-	sections := make([]map[string]any, 0, len(enriched))
+	sections := make([]map[string]any, 0, len(pages)+len(enriched)+len(runnerUpPages))
 	finalIDs := make([]tree.SectionID, 0, len(enriched))
 	// A page-based strategy's evidence pages lead the list, as-is: the
 	// page the Judge found holds the answer; a section mapped to it by
-	// the parser's page attribution may not (HAL-1390).
-	if selResult != nil {
-		sections = append(sections, evidencePageSections(selResult.EvidencePages)...)
-	}
+	// the parser's page attribution may not (HAL-1390). Pages ranked just
+	// below the evidence close the list when there was room for them.
+	sections = append(sections, evidencePageSections(pages)...)
 	for _, e := range enriched {
 		sections = append(sections, sectionWithContentToMap(e))
 		finalIDs = append(finalIDs, e.sec.ID)
 	}
+	for _, m := range evidencePageSections(runnerUpPages) {
+		m["runner_up"] = true
+		sections = append(sections, m)
+	}
 
 	// Trace token is computed over the FINAL IDs that ship in the
 	// response (after max_sections + ReRank truncation). Recomputing
@@ -709,7 +699,8 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
 		"trace_token": traceToken,
 		// What retrieval cost. /v1/answer always reported this; /v1/query
 		// dropped it, so a caller benchmarking retrieval alone saw $0.
-		"usage": usageJSON(totalUsage),
+		"usage":      usageJSON(totalUsage),
+		"unit_stats": unitStats,
 	}
 	if plan != nil {
 		resp["plan"] = plan
@@ -740,6 +731,41 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
+// assembleUnits builds the units a response returns (retrieval.AssembleUnits)
+// and splits them the way the handlers use them: evidence pages, sections
+// with their text (re-rank and span extraction work on these), and pages
+// ranked below the evidence that fill the slots left over.
+func (d Deps) assembleUnits(ctx context.Context, t *tree.Tree, res *retrieval.Result, k int) ([]retrieval.EvidencePage, []sectionWithContent, []retrieval.EvidencePage, retrieval.UnitStats) {
+	units, stats := retrieval.AssembleUnits(ctx, t, res, k, d.loadContent)
+	var pages, runnerUps []retrieval.EvidencePage
+	var secs []sectionWithContent
+	for _, u := range units {
+		switch {
+		case u.Kind == retrieval.UnitPage && u.RunnerUp:
+			runnerUps = append(runnerUps, u.Page)
+		case u.Kind == retrieval.UnitPage:
+			pages = append(pages, u.Page)
+		default:
+			secs = append(secs, sectionWithContent{sec: u.Section, content: u.Content})
+		}
+	}
+	return pages, secs, runnerUps, stats
+}
+
+// loadContent reads a section's stored text.
+func (d Deps) loadContent(ctx context.Context, ref string) (string, error) {
+	rc, _, err := d.Storage.Get(ctx, ref)
+	if err != nil {
+		return "", err
+	}
+	defer rc.Close() // best-effort close
+	raw, err := io.ReadAll(rc)
+	if err != nil {
+		return "", err
+	}
+	return string(raw), nil
+}
+
 // sectionWithContent bundles a tree section with its loaded content
 // and optional re-rank score / answer-span. Used by /v1/query and
 // /v1/answer.
@@ -946,28 +972,9 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
 	if maxSections <= 0 {
 		maxSections = 5
 	}
-	if len(ids) > maxSections {
-		ids = ids[:maxSections]
-	}
-
-	// Load each section's content.
-	enriched := make([]sectionWithContent, 0, len(ids))
-	for _, id := range ids {
-		sec := t.FindByID(id)
-		if sec == nil {
-			continue
-		}
-		var content string
-		if sec.ContentRef != "" {
-			rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef)
-			if err == nil {
-				raw, _ := io.ReadAll(rc)
-				_ = rc.Close() // best-effort close
-				content = string(raw)
-			}
-		}
-		enriched = append(enriched, sectionWithContent{sec: sec, content: content})
-	}
+	// Each section the answer is written from carries body text
+	// (HAL-3160).
+	_, enriched, _, _ := d.assembleUnits(r.Context(), t, &retrieval.Result{SelectedIDs: ids}, maxSections)
 
 	// Optional: content-aware re-rank before synthesis sees the
 	// evidence. When TopK is set the synthesis prompt only ever sees
@@ -1194,20 +1201,9 @@ func (d Deps) handleQueryMulti(w http.ResponseWriter, r *http.Request) {
 	docs := make([]map[string]any, 0, len(result.Documents))
 	for docID, dr := range result.Documents {
 		sections := make([]map[string]any, 0, len(dr.SelectedIDs))
-		for _, sid := range dr.SelectedIDs {
-			sec := dr.Tree.FindByID(sid)
-			if sec == nil {
-				continue
-			}
-			var content string
-			if sec.ContentRef != "" {
-				rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef)
-				if err == nil {
-					raw, _ := io.ReadAll(rc)
-					_ = rc.Close() // best-effort close
-					content = string(raw)
-				}
-			}
+		_, secs, _, _ := d.assembleUnits(r.Context(), dr.Tree, &retrieval.Result{SelectedIDs: dr.SelectedIDs}, body.MaxSections)
+		for _, e := range secs {
+			sec, content := e.sec, e.content
 			s := map[string]any{
 				"id":          sec.ID,
 				"parent_id":   sec.ParentID,
@@ -1226,9 +1222,6 @@ func (d Deps) handleQueryMulti(w http.ResponseWriter, r *http.Request) {
 				s["candidate_questions"] = sec.CandidateQuestions
 			}
 			sections = append(sections, s)
-			if body.MaxSections > 0 && len(sections) >= body.MaxSections {
-				break
-			}
 		}
 		docs = append(docs, map[string]any{
 			"document_id": docID,

internal/handler/answer.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"log/slog"
 	"net/http"
 	"strings"
@@ -341,26 +340,14 @@ func (h *AnswerHandler) answerFromSections(ctx context.Context, t *tree.Tree, bo
 	if maxSections <= 0 {
 		maxSections = 5
 	}
-	if len(ids) > maxSections {
-		ids = ids[:maxSections]
-	}
-
-	enriched := make([]answerSection, 0, len(ids))
-	for _, id := range ids {
-		sec := t.FindByID(id)
-		if sec == nil {
-			continue
-		}
-		var content string
-		if sec.ContentRef != "" {
-			rc, _, getErr := h.storage.Get(ctx, sec.ContentRef)
-			if getErr == nil {
-				raw, _ := io.ReadAll(rc)
-				_ = rc.Close() // best-effort close
-				content = string(raw)
-			}
-		}
-		enriched = append(enriched, answerSection{sec: sec, content: content})
+	// Each section the answer is written from carries body text; a
+	// heading-only one, or a repeat, gives its slot to the next-ranked
+	// leaf (HAL-3160).
+	units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{SelectedIDs: ids, RunnerUpIDs: res.RunnerUpIDs}, maxSections,
+		func(ctx context.Context, ref string) (string, error) { return readStored(ctx, h.storage, ref) })
+	enriched := make([]answerSection, 0, len(units))
+	for _, u := range units {
+		enriched = append(enriched, answerSection{sec: u.Section, content: u.Content})
 	}
 
 	spanExtractor := h.spanExtractor(body.Model)

internal/handler/answer_pages.go

@@ -54,9 +54,14 @@ func (h *AnswerHandler) answerFromPages(
 	if k <= 0 {
 		k = defaultEvidencePages
 	}
-	evidence := res.EvidencePages
-	if len(evidence) > k {
-		evidence = evidence[:k]
+	// The answer reads the evidence pages and, when the navigator found
+	// fewer than k above its threshold, the pages it ranked next, so a
+	// question that needs two statements is not answered from one
+	// (HAL-3160).
+	units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{EvidencePages: res.EvidencePages, RunnerUpPages: res.RunnerUpPages}, k, nil)
+	evidence := make([]retrieval.EvidencePage, 0, len(units))
+	for _, u := range units {
+		evidence = append(evidence, u.Page)
 	}
 	retrievalMS := time.Since(started).Milliseconds()
 

internal/handler/query.go

@@ -161,6 +161,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) {
 		ids      []tree.SectionID
 		usage    *retrieval.Usage
 		evidence []retrieval.EvidencePage
+		full     *retrieval.Result
 	)
 	if cs, ok := strategy.(retrieval.CostStrategy); ok {
 		result, err := cs.SelectWithCost(r.Context(), t, body.Query, budget)
@@ -175,6 +176,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) {
 		ids = result.SelectedIDs
 		usage = &result.Usage
 		evidence = result.EvidencePages
+		full = result
 	} else {
 		var err error
 		ids, err = strategy.Select(r.Context(), t, body.Query, budget)
@@ -192,48 +194,23 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) {
 	// page the Judge chose holds the answer, while the sections overlapping
 	// it come back in document order (HAL-1390). max_sections caps the
 	// combined list, so a small cap keeps the evidence, not the first
-	// sections of the document (HAL-3071).
-	sections := evidencePageSections(evidence)
-	if body.MaxSections > 0 {
-		if len(sections) > body.MaxSections {
-			sections = sections[:body.MaxSections]
-		}
-		if room := body.MaxSections - len(sections); len(ids) > room {
-			ids = ids[:room]
-		}
-	}
-
-	for _, id := range ids {
-		sec := t.FindByID(id)
-		if sec == nil {
-			continue
-		}
-		var content string
-		if sec.ContentRef != "" {
-			rc, _, getErr := h.storage.Get(r.Context(), sec.ContentRef)
-			if getErr == nil {
-				raw, _ := io.ReadAll(rc)
-				_ = rc.Close() // best-effort close
-				content = string(raw)
-			}
-		}
-		sections = append(sections, map[string]any{
-			"id":          sec.ID,
-			"parent_id":   sec.ParentID,
-			"title":       sec.Title,
-			"summary":     sec.Summary,
-			"token_count": sec.TokenCount,
-			"content":     content,
-		})
+	// sections of the document (HAL-3071). Every unit carries body text: a
+	// heading-only section, or one the pages already carry, gives its slot
+	// to the next-ranked evidence (HAL-3160).
+	res := &retrieval.Result{SelectedIDs: ids, EvidencePages: evidence}
+	if full != nil {
+		res = full
 	}
-
+	units, unitStats := retrieval.AssembleUnits(r.Context(), t, res, body.MaxSections, h.loadContent)
+	sections := unitMaps(units)
 	resp := map[string]any{
 		"document_id": body.DocumentID,
 		"query":       body.Query,
 		"strategy":    retrieval.PublicName(strategy.Name()),
 		"model":       body.Model,
 		"sections":    sections,
 		"elapsed_ms":  time.Since(started).Milliseconds(),
+		"unit_stats":  unitStats,
 	}
 	if usage != nil {
 		resp["usage"] = map[string]any{
@@ -249,6 +226,63 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) {
 	writeJSON(w, http.StatusOK, resp)
 }
 
+// loadContent reads a section's stored text.
+func (h *QueryHandler) loadContent(ctx context.Context, ref string) (string, error) {
+	return readStored(ctx, h.storage, ref)
+}
+
+// readStored reads one stored object as text.
+func readStored(ctx context.Context, store storage.Storage, ref string) (string, error) {
+	rc, _, err := store.Get(ctx, ref)
+	if err != nil {
+		return "", err
+	}
+	defer rc.Close() // best-effort close
+	raw, err := io.ReadAll(rc)
+	if err != nil {
+		return "", err
+	}
+	return string(raw), nil
+}
+
+// unitMaps renders assembled units in the response shape: a page as
+// evidencePageSections renders it, a section with its tree fields. A unit
+// that fills a slot below the evidence says so (runner_up), and a heading
+// whose text was gathered from the sections beneath it says that
+// (expanded).
+func unitMaps(units []retrieval.Unit) []map[string]any {
+	out := make([]map[string]any, 0, len(units))
+	for _, u := range units {
+		var m map[string]any
+		if u.Kind == retrieval.UnitPage {
+			m = evidencePageSections([]retrieval.EvidencePage{u.Page})[0]
+		} else {
+			m = map[string]any{
+				"id":          u.Section.ID,
+				"parent_id":   u.Section.ParentID,
+				"title":       u.Section.Title,
+				"summary":     u.Section.Summary,
+				"token_count": u.Section.TokenCount,
+				"content":     u.Content,
+			}
+			if u.Section.PageStart > 0 {
+				m["page_start"] = u.Section.PageStart
+			}
+			if u.Section.PageEnd > 0 {
+				m["page_end"] = u.Section.PageEnd
+			}
+			if u.Expanded {
+				m["expanded"] = true
+			}
+		}
+		if u.RunnerUp {
+			m["runner_up"] = true
+		}
+		out = append(out, m)
+	}
+	return out
+}
+
 // evidencePageSections turns a page-based strategy's evidence pages into
 // response units: id "page_<n>", the owning section's title with the page
 // ("<section> · p.<n>"), the page text as content. The same shape the