feat(search): filter before ranking and expose citation handles

src/search.test.ts

@@ -0,0 +1,90 @@
+import { describe, expect, it } from 'vitest'
+import { type SearchKnowledgeOptions, searchKnowledge } from './search'
+import type { KnowledgeIndex, KnowledgePage } from './types'
+
+function page(
+  id: string,
+  title: string,
+  text: string,
+  options: { kind?: string; tags?: string[]; outLinks?: string[]; sourceIds?: string[] } = {},
+): KnowledgePage {
+  return {
+    id,
+    path: `knowledge/${id}.md`,
+    title,
+    text,
+    frontmatter: options.kind ? { kind: options.kind } : {},
+    sourceIds: options.sourceIds ?? [],
+    tags: options.tags ?? [],
+    outLinks: options.outLinks ?? [],
+  }
+}
+
+function index(pages: KnowledgePage[]): KnowledgeIndex {
+  return {
+    root: '/kb',
+    generatedAt: '2026-08-17T00:00:00.000Z',
+    sources: [],
+    pages,
+    graph: { nodes: [], edges: [] },
+  }
+}
+
+describe('searchKnowledge filters', () => {
+  const pages = [
+    page('prior-alpha', 'Alpha exact prior', 'alpha alpha alpha', {
+      kind: 'prior',
+      tags: ['curated', 'math'],
+    }),
+    page('finding-alpha', 'Alpha measured finding', 'alpha measurement', {
+      kind: 'finding',
+      tags: ['measured', 'math'],
+    }),
+    page('profile-alpha', 'Alpha profile', 'alpha track record', {
+      kind: 'profile',
+      tags: ['agent'],
+    }),
+  ]
+
+  it('keeps the numeric limit overload and returns an explicit citation handle', () => {
+    const hits = searchKnowledge(index(pages), 'alpha', 1)
+
+    expect(hits).toHaveLength(1)
+    expect(hits[0]?.citationId).toBe(hits[0]?.page.id)
+    expect(hits[0]?.rank).toBe(1)
+  })
+
+  it('applies kind and tag filters before either ranking stage', () => {
+    const hits = searchKnowledge(index(pages), 'alpha', {
+      kinds: ['finding'],
+      tags: ['measured'],
+    })
+
+    expect(hits.map((hit) => hit.citationId)).toEqual(['finding-alpha'])
+    expect(hits[0]?.normalizedScore).toBe(1)
+  })
+
+  it('combines stable ids with a caller-owned predicate', () => {
+    const options: SearchKnowledgeOptions = {
+      pageIds: ['prior-alpha', 'finding-alpha'],
+      predicate: (candidate) => candidate.title.includes('prior'),
+    }
+
+    expect(searchKnowledge(index(pages), 'alpha', options).map((hit) => hit.citationId)).toEqual([
+      'prior-alpha',
+    ])
+  })
+
+  it('treats an explicit empty filter as matching no pages', () => {
+    expect(searchKnowledge(index(pages), 'alpha', { kinds: [] })).toEqual([])
+    expect(searchKnowledge(index(pages), 'alpha', { tags: [] })).toEqual([])
+    expect(searchKnowledge(index(pages), 'alpha', { pageIds: [] })).toEqual([])
+  })
+
+  it('refuses a malformed limit instead of changing its meaning', () => {
+    expect(() => searchKnowledge(index(pages), 'alpha', -1)).toThrow(/non-negative integer/)
+    expect(() => searchKnowledge(index(pages), 'alpha', { limit: 1.5 })).toThrow(
+      /non-negative integer/,
+    )
+  })
+})

src/search.ts

@@ -1,4 +1,4 @@
-import type { KnowledgeIndex, KnowledgePage, KnowledgeSearchResult } from './types'
+import type { KnowledgeId, KnowledgeIndex, KnowledgePage, KnowledgeSearchResult } from './types'
 
 const RRF_K = 60
 const STOP_WORDS = new Set([
@@ -21,17 +21,59 @@ const STOP_WORDS = new Set([
   'and',
 ])
 
+export interface SearchKnowledgeOptions {
+  /** Maximum results returned. Defaults to 10. */
+  limit?: number
+  /** Match pages whose stable id is one of these values. */
+  pageIds?: readonly KnowledgeId[]
+  /** Match pages carrying at least one of these tags. */
+  tags?: readonly string[]
+  /** Match the exact string stored in `frontmatter.kind`. */
+  kinds?: readonly string[]
+  /** Additional caller-owned filter, applied before either ranking stage. */
+  predicate?: (page: KnowledgePage) => boolean
+}
+
+/**
+ * A retrieval result with an explicit citation handle.
+ *
+ * `citationId` is exactly `page.id`; later writes should persist this value
+ * when they cite the page. Keeping it at the result's top level prevents tool
+ * renderers from accidentally hiding the only stable handle a model can copy.
+ */
+export interface KnowledgeSearchHit extends KnowledgeSearchResult {
+  citationId: KnowledgeId
+}
+
+export function searchKnowledge(
+  index: KnowledgeIndex,
+  query: string,
+  limit?: number,
+): KnowledgeSearchHit[]
+export function searchKnowledge(
+  index: KnowledgeIndex,
+  query: string,
+  options?: SearchKnowledgeOptions,
+): KnowledgeSearchHit[]
 export function searchKnowledge(
   index: KnowledgeIndex,
   query: string,
-  limit = 10,
-): KnowledgeSearchResult[] {
+  limitOrOptions: number | SearchKnowledgeOptions = 10,
+): KnowledgeSearchHit[] {
   const trimmed = query.trim()
   if (trimmed === '') return []
-  const tokenRanked = rankByTokens(index.pages, trimmed)
-  const graphRanked = rankByGraph(index.pages, tokenRanked)
+  const options =
+    typeof limitOrOptions === 'number' ? { limit: limitOrOptions } : { ...limitOrOptions }
+  const limit = options.limit ?? 10
+  if (!Number.isInteger(limit) || limit < 0) {
+    throw new Error(`search limit must be a non-negative integer, got ${String(limit)}`)
+  }
+
+  const pages = filterPages(index.pages, options)
+  const tokenRanked = rankByTokens(pages, trimmed)
+  const graphRanked = rankByGraph(pages, tokenRanked)
   const scores = reciprocalRankFusion([tokenRanked.map((p) => p.id), graphRanked.map((p) => p.id)])
-  const byId = new Map(index.pages.map((page) => [page.id, page]))
+  const byId = new Map(pages.map((page) => [page.id, page]))
 
   const ranked = [...scores.entries()]
     .map(([id, score]) => ({ page: byId.get(id), score }))
@@ -46,6 +88,7 @@ export function searchKnowledge(
   const topScore = ranked[0]?.score ?? 0
 
   return ranked.map((item, i) => ({
+    citationId: item.page.id,
     page: item.page,
     score: item.score,
     rrfScore: item.score,
@@ -83,6 +126,22 @@ export function reciprocalRankFusion(rankLists: string[][], k = RRF_K): Map<stri
   return scores
 }
 
+function filterPages(pages: KnowledgePage[], options: SearchKnowledgeOptions): KnowledgePage[] {
+  const pageIds = options.pageIds ? new Set(options.pageIds) : null
+  const tags = options.tags ? new Set(options.tags) : null
+  const kinds = options.kinds ? new Set(options.kinds) : null
+
+  return pages.filter((page) => {
+    if (pageIds && !pageIds.has(page.id)) return false
+    if (tags && !page.tags.some((tag) => tags.has(tag))) return false
+    if (kinds) {
+      const kind = page.frontmatter.kind
+      if (typeof kind !== 'string' || !kinds.has(kind)) return false
+    }
+    return options.predicate?.(page) ?? true
+  })
+}
+
 function rankByTokens(pages: KnowledgePage[], query: string): KnowledgePage[] {
   const tokens = tokenizeQuery(query)
   const effective = tokens.length > 0 ? tokens : [query.toLowerCase()]