CPLAT-11801: show matched role in search results and rank user hits first

README.md

@@ -190,7 +190,11 @@ Search inside conversation content across all sessions (`Ctrl+S` or `:search`).
 
 **Features:**
 - Searches text, tool inputs, thinking blocks, and system tags
-- Results are ordered newest session first
+- Each hit is prefixed with the role that matched (` usr` / ` ast`), so a hit in
+  your own prompt is distinguishable from one in the model's reply
+- Results are ordered newest session first, and **within a session the user's own
+  words come first** — a hit in your prompt is usually the one you were looking
+  for. Session recency still wins across sessions, so the list stays browsable
 - Matched terms are highlighted in snippets
 - `[LIVE]` / `[HERE]` badges show which hits belong to a running session, updated
   while the results are open

internal/session/index.go

@@ -545,6 +545,7 @@ type indexHit struct {
 	path     string
 	lineOff  int64
 	blockIdx int
+	role     string
 }
 
 // queryIndex returns matching block locations, newest session first.
@@ -574,17 +575,27 @@ func (ix *Index) queryIndex(ctx context.Context, q SearchQuery, allowed map[stri
 		}
 	}
 
-	sqlText := `select files.path, locs.line_off, locs.block_idx
+	sqlText := `select files.path, locs.line_off, locs.block_idx, locs.role
 	            from blocks
 	            join locs  on locs.rowid = blocks.rowid
 	            join files on files.id   = locs.file_id
 	            where blocks match ?`
 	if len(where) > 0 {
 		sqlText += " and " + strings.Join(where, " and ")
 	}
-	// Ordering is applied in Go against session mtime; SQL only bounds the work.
+	// Bound the work with a role-ordered cut, sized so session ranking still has
+	// room.
+	//
+	// Session order needs session mtime, which the index does not store, so it
+	// happens in Go. A plain SQL LIMIT truncates in rowid order and can drop
+	// rows the Go sort would have kept; ordering by role alone is global and
+	// starves recent sessions. Over-fetching a multiple of the limit under a
+	// role ordering gives both: the user's words survive the cut, and enough
+	// rows from enough files remain for the recency sort to be meaningful.
 	if limit > 0 {
-		sqlText += fmt.Sprintf(" limit %d", limit)
+		sqlText += ` order by case locs.role
+		                      when 'user' then 0 when 'assistant' then 1 else 2 end`
+		sqlText += fmt.Sprintf(" limit %d", limit*sqlOverFetch)
 	}
 
 	rows, err := ix.db.QueryContext(ctx, sqlText, args...)
@@ -596,7 +607,7 @@ func (ix *Index) queryIndex(ctx context.Context, q SearchQuery, allowed map[stri
 	var hits []indexHit
 	for rows.Next() {
 		var h indexHit
-		if err := rows.Scan(&h.path, &h.lineOff, &h.blockIdx); err != nil {
+		if err := rows.Scan(&h.path, &h.lineOff, &h.blockIdx, &h.role); err != nil {
 			return nil, err
 		}
 		if allowed != nil {
@@ -611,7 +622,14 @@ func (ix *Index) queryIndex(ctx context.Context, q SearchQuery, allowed map[stri
 	}
 
 	// Newest session first, matching the session browser's ordering; within a
-	// session, transcript order.
+	// session, what you said comes before what the model said, then transcript
+	// order.
+	//
+	// The user-first tie-break is deliberately INSIDE the session, not across
+	// them: hoisting every user hit to the top would break the recency ordering
+	// that makes "what was I doing lately" browsable. Ordering here rather than
+	// after hydration also means a limit truncates assistant hits first, so a
+	// broad query cannot fill its cap with replies and drop the prompts.
 	sort.SliceStable(hits, func(i, j int) bool {
 		si, sj := allowed[hits[i].path], allowed[hits[j].path]
 		if si != nil && sj != nil && !si.ModTime.Equal(sj.ModTime) {
@@ -620,7 +638,38 @@ func (ix *Index) queryIndex(ctx context.Context, q SearchQuery, allowed map[stri
 		if hits[i].path != hits[j].path {
 			return hits[i].path < hits[j].path
 		}
+		if ri, rj := roleRank(hits[i].role), roleRank(hits[j].role); ri != rj {
+			return ri < rj
+		}
 		return hits[i].lineOff < hits[j].lineOff
 	})
 	return hits, nil
 }
+
+// roleRank orders matched entries within one session: the user's own words
+// first. A hit in your prompt is usually the one you were looking for — it is
+// what you wrote and therefore what you remember — while the model's reply is
+// the elaboration around it. Anything without a role sorts last rather than
+// being mixed into either group.
+func roleRank(role string) int {
+	switch role {
+	case "user":
+		return 0
+	case "assistant":
+		return 1
+	default:
+		return 2
+	}
+}
+
+// sqlOverFetch is how many times the caller's limit is pulled from SQL before
+// Go ranks and trims. The SQL cut can only order by role (session mtime lives
+// outside the index), so the surplus is what keeps the recency sort from being
+// decided by an arbitrary truncation.
+//
+// It is a bound, not a guarantee: a query matching more than limit*sqlOverFetch
+// user blocks inside a single old session could still crowd out a recent one.
+// Raising it trades latency for that tail — measured on a 400 MB index, a very
+// broad query costs ~150ms at 8x and ~400ms unbounded, against ~1.8s for the
+// full scan this replaced.
+const sqlOverFetch = 8

internal/session/search_index.go

@@ -51,14 +51,12 @@ func SearchWithIndex(ctx context.Context, ix *Index, sessions []*Session, q Sear
 		}
 	}
 
-	// Over-fetch: post-filtering (exclusions, tool prefix) drops some hits, and
-	// a limit applied in SQL would otherwise silently truncate good results.
-	sqlLimit := 0
-	if limit > 0 {
-		sqlLimit = limit * 4
-	}
-
-	hits, err := ix.queryIndex(ctx, q, bySession, sqlLimit)
+	// The limit is applied per FILE inside queryIndex, not globally: session
+	// ranking happens here in Go (it needs session mtime), so a global cut would
+	// discard rows the ranking would have kept. Every file contributing its own
+	// best `limit` rows is enough for any ordering this function can produce,
+	// while keeping a broad query from dragging back tens of thousands of rows.
+	hits, err := ix.queryIndex(ctx, q, bySession, limit)
 	if err != nil {
 		// An index failure is not a search failure — degrade to the scan.
 		return collectScan(ctx, sessions, q, limit), SearchModeScan, nil
@@ -206,15 +204,31 @@ func collectScan(ctx context.Context, sessions []*Session, q SearchQuery, limit
 		}
 	}
 
+	// Same order as the indexed path: newest session first, and within one
+	// session the user's own words before the model's (see queryIndex).
 	sort.SliceStable(out, func(i, j int) bool {
 		si, sj := out[i].Session, out[j].Session
 		if si == nil || sj == nil {
 			return false
 		}
-		return si.ModTime.After(sj.ModTime)
+		if !si.ModTime.Equal(sj.ModTime) {
+			return si.ModTime.After(sj.ModTime)
+		}
+		if si.ID != sj.ID {
+			return si.ID < sj.ID
+		}
+		return resultRoleRank(out[i]) < resultRoleRank(out[j])
 	})
 	if limit > 0 && len(out) > limit {
 		out = out[:limit]
 	}
 	return out
 }
+
+// resultRoleRank is roleRank for a hydrated result.
+func resultRoleRank(r SearchResult) int {
+	if r.Entry == nil {
+		return 2
+	}
+	return roleRank(r.Entry.Role)
+}

internal/session/search_role_test.go

@@ -0,0 +1,199 @@
+package session
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+)
+
+// writeRoleTranscript writes a transcript whose lines alternate by role, each
+// containing the marker so every entry is a hit.
+func writeRoleTranscript(t *testing.T, dir, name string, roles []string, marker string) *Session {
+	t.Helper()
+	var lines []string
+	for i, role := range roles {
+		text := fmt.Sprintf("%s occurrence %d", marker, i)
+		lines = append(lines, fmt.Sprintf(
+			`{"type":%q,"uuid":"u%d","message":{"role":%q,"content":[{"type":"text","text":%q}]}}`,
+			role, i, role, text))
+	}
+	path := filepath.Join(dir, name)
+	if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")+"\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	fi, err := os.Stat(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return &Session{ID: name, ShortID: name, FilePath: path, ModTime: fi.ModTime(), ProjectName: name}
+}
+
+func resultRoles(rs []SearchResult) []string {
+	out := make([]string, 0, len(rs))
+	for _, r := range rs {
+		if r.Entry == nil {
+			out = append(out, "")
+			continue
+		}
+		out = append(out, r.Entry.Role)
+	}
+	return out
+}
+
+// Within one session the user's own words come first: a hit in your prompt is
+// what you wrote, and therefore what you remember searching for.
+func TestUserHitsRankBeforeAssistantWithinSession(t *testing.T) {
+	ix, dir := openTestIndex(t)
+	s := writeRoleTranscript(t, dir, "a.jsonl",
+		[]string{"assistant", "user", "assistant", "user"}, "marker")
+	syncAll(t, ix, []*Session{s})
+
+	res, _ := searchIdx(t, ix, []*Session{s}, "marker")
+	if len(res) != 4 {
+		t.Fatalf("hits = %d, want 4", len(res))
+	}
+	got := resultRoles(res)
+	want := []string{"user", "user", "assistant", "assistant"}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("role order = %v, want %v", got, want)
+		}
+	}
+}
+
+// User-first must NOT reorder sessions: recency is what makes the result list
+// browsable, and hoisting every user hit across sessions would destroy it.
+func TestUserPriorityDoesNotReorderSessions(t *testing.T) {
+	ix, dir := openTestIndex(t)
+	// Older session whose hits are all user; newer session whose hits are all
+	// assistant. Recency must still win at the session level.
+	older := writeRoleTranscript(t, dir, "older.jsonl", []string{"user", "user"}, "marker")
+	newer := writeRoleTranscript(t, dir, "newer.jsonl", []string{"assistant", "assistant"}, "marker")
+
+	past := time.Now().Add(-48 * time.Hour)
+	os.Chtimes(older.FilePath, past, past)
+	fi, _ := os.Stat(older.FilePath)
+	older.ModTime = fi.ModTime()
+
+	sessions := []*Session{older, newer}
+	syncAll(t, ix, sessions)
+
+	res, _ := searchIdx(t, ix, sessions, "marker")
+	if len(res) != 4 {
+		t.Fatalf("hits = %d, want 4", len(res))
+	}
+	// The newer (assistant-only) session must still come first.
+	if res[0].Session.FilePath != newer.FilePath {
+		t.Errorf("first hit is from %s, want the newer session",
+			filepath.Base(res[0].Session.FilePath))
+	}
+	if res[len(res)-1].Session.FilePath != older.FilePath {
+		t.Errorf("last hit is from %s, want the older session",
+			filepath.Base(res[len(res)-1].Session.FilePath))
+	}
+}
+
+// The limit must not be filled by assistant hits while user hits are dropped —
+// which is why the ordering happens before truncation, not after.
+func TestLimitKeepsUserHitsOverAssistant(t *testing.T) {
+	ix, dir := openTestIndex(t)
+	roles := make([]string, 0, 12)
+	for i := 0; i < 10; i++ {
+		roles = append(roles, "assistant")
+	}
+	roles = append(roles, "user", "user")
+	s := writeRoleTranscript(t, dir, "a.jsonl", roles, "marker")
+	syncAll(t, ix, []*Session{s})
+
+	res, _, err := SearchWithIndex(context.Background(), ix, []*Session{s}, ParseSearchQuery("marker"), 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res) != 2 {
+		t.Fatalf("hits = %d, want 2", len(res))
+	}
+	for i, r := range res {
+		if r.Entry == nil || r.Entry.Role != "user" {
+			t.Errorf("result %d role = %v, want user (assistant hits filled the limit)",
+				i, resultRoles(res))
+			break
+		}
+	}
+}
+
+// The scan fallback must order the same way, or results reshuffle depending on
+// whether the index could answer the query.
+func TestScanFallbackUsesSameRoleOrdering(t *testing.T) {
+	dir := t.TempDir()
+	s := writeRoleTranscript(t, dir, "a.jsonl",
+		[]string{"assistant", "user", "assistant", "user"}, "marker")
+
+	scan := collectScan(context.Background(), []*Session{s}, ParseSearchQuery("marker"), 0)
+	if len(scan) != 4 {
+		t.Fatalf("scan hits = %d, want 4", len(scan))
+	}
+	got := resultRoles(scan)
+	want := []string{"user", "user", "assistant", "assistant"}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("scan role order = %v, want %v", got, want)
+		}
+	}
+}
+
+func TestRoleRankOrdersUserFirst(t *testing.T) {
+	if roleRank("user") >= roleRank("assistant") {
+		t.Error("user must rank before assistant")
+	}
+	if roleRank("assistant") >= roleRank("") {
+		t.Error("a roleless entry must sort after both, not among them")
+	}
+}
+
+// Ranking must never starve a recent session. An earlier attempt ordered by
+// role in SQL so the LIMIT would keep user hits; that ranking is global, so an
+// old session with many user hits filled the cap and the newest session
+// vanished from the results entirely. This pins the failure that caught it.
+func TestLimitDoesNotStarveRecentSessions(t *testing.T) {
+	ix, dir := openTestIndex(t)
+	oldRoles := make([]string, 0, 20)
+	for i := 0; i < 20; i++ {
+		oldRoles = append(oldRoles, "user")
+	}
+	older := writeRoleTranscript(t, dir, "older.jsonl", oldRoles, "marker")
+	// The recent session's user hit sits DEEP in its file, so a SQL order of
+	// (role, line_off) would rank it behind the old session's shallow hits.
+	newRoles := make([]string, 0, 40)
+	for i := 0; i < 39; i++ {
+		newRoles = append(newRoles, "assistant")
+	}
+	newRoles = append(newRoles, "user")
+	newer := writeRoleTranscript(t, dir, "newer.jsonl", newRoles, "marker")
+
+	past := time.Now().Add(-72 * time.Hour)
+	os.Chtimes(older.FilePath, past, past)
+	fi, _ := os.Stat(older.FilePath)
+	older.ModTime = fi.ModTime()
+
+	sessions := []*Session{older, newer}
+	syncAll(t, ix, sessions)
+
+	res, _, err := SearchWithIndex(context.Background(), ix, sessions, ParseSearchQuery("marker"), 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var sawNewer bool
+	for _, r := range res {
+		if r.Session.FilePath == newer.FilePath {
+			sawNewer = true
+		}
+	}
+	t.Logf("results=%d sawNewer=%v", len(res), sawNewer)
+	if !sawNewer {
+		t.Error("the most recent session is absent from a limited result set")
+	}
+}