Files
Plateau Nguyenandviettranx 6d7389539a fix(vault): prevent vault_read id-namespace collision (#959)
* fix(vault): prevent vault_read id-namespace collision

vault_search was leaking KG/episodic entity ids into result sets even when
narrow `types` were requested, and callers then passed those ids to
vault_read which returned a generic "document not found". The cause was
threefold:

1. `types` filter was only applied to the vault fan-out; KG and episodic
   ran unconditionally. Now gated by shouldFanout(types, key).
2. vault_search output lacked a per-source tool hint. Each result now ends
   with " → use <tool>" naming the correct follow-up (vault_read,
   knowledge_graph_search, or memory_search).
3. vault_read miss returned "document not found" without checking whether
   the id belonged to a foreign namespace. It now probes KG then episodic
   and returns a namespace-specific redirect error. Stores are injected
   via SetKGStore/SetEpisodicStore, nil-safe, tenant-scoped.

Adds red→green characterization tests plus an end-to-end integration
scenario seeding a vault doc + KG entity with identical basenames.

* test(agent): bump none-mode prompt size budget to 3100

vault_read wiring (#948) added ~95 chars to read_file tool summary,
pushing none-mode prompt from <3000 to 3075 chars. Bump budget to
3100 (~775 tokens) to match the intentional addition.

* test(integration): ensure data_migrations table exists in reset helper

The reset helper runs before RunPendingHooks, but RunPendingHooks is
what normally creates data_migrations. On a fresh CI database the
DELETE fails with 'relation does not exist'. Create the table
defensively so reset works regardless of execution order.

* refactor(vault): per-source id fields + wire episodic into search

Align vault_search output fields with downstream tool input params:
doc_id (vault_read), entity_id (knowledge_graph_search), episodic_id
(memory_expand). Prevents LLMs from pattern-matching a generic `id:`
and misrouting a foreign-namespace uuid into vault_read. Fallback
redirect in vault_read now quotes id + names the correct param so
the LLM can self-correct in one turn.

Also wire stores.Episodic into VaultSearchService (stale comment
claimed pending-impl; PGEpisodicStore has existed and been in use
since v3). Unifies search fan-out with vault_read namespace probe.

---------

Co-authored-by: viettranx <viettranx@gmail.com>
2026-04-19 15:00:07 +07:00

160 lines
4.5 KiB
Go

package vault
import (
"context"
"testing"
"time"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// fakeVaultStoreSearch embeds store.VaultStore to satisfy the interface;
// only Search is exercised.
type fakeVaultStoreSearch struct {
store.VaultStore
results []store.VaultSearchResult
calls int
}
func (f *fakeVaultStoreSearch) Search(ctx context.Context, opts store.VaultSearchOptions) ([]store.VaultSearchResult, error) {
f.calls++
return f.results, nil
}
type fakeEpisodicStoreSearch struct {
store.EpisodicStore
results []store.EpisodicSearchResult
calls int
}
func (f *fakeEpisodicStoreSearch) Search(ctx context.Context, query string, agentID, userID string, opts store.EpisodicSearchOptions) ([]store.EpisodicSearchResult, error) {
f.calls++
return f.results, nil
}
type fakeKGStoreSearch struct {
store.KnowledgeGraphStore
entities []store.Entity
calls int
}
func (f *fakeKGStoreSearch) SearchEntities(ctx context.Context, agentID, userID, query string, limit int) ([]store.Entity, error) {
f.calls++
return f.entities, nil
}
func makeVaultResult(id, title string) store.VaultSearchResult {
return store.VaultSearchResult{
Document: store.VaultDocument{ID: id, Title: title, Path: title + ".md", DocType: "context"},
Score: 0.8,
Source: "vault",
}
}
func makeEpisodicResult(id, key string) store.EpisodicSearchResult {
return store.EpisodicSearchResult{
EpisodicID: id,
SessionKey: key,
L0Abstract: "abstract",
Score: 0.5,
CreatedAt: time.Now(),
}
}
func makeEntity(id, name string) store.Entity {
return store.Entity{ID: id, Name: name, EntityType: "document", Confidence: 0.7}
}
// --- Test 1: types filter excludes KG/episodic when a narrow type is requested. ---
func TestSearch_TypesFilterExcludesKG(t *testing.T) {
vs := &fakeVaultStoreSearch{results: []store.VaultSearchResult{makeVaultResult("V1", "VaultDoc")}}
es := &fakeEpisodicStoreSearch{results: []store.EpisodicSearchResult{makeEpisodicResult("E1", "sess")}}
kg := &fakeKGStoreSearch{entities: []store.Entity{makeEntity("K1", "KGDoc")}}
svc := NewVaultSearchService(vs, es, kg)
res, err := svc.Search(context.Background(), UnifiedSearchOptions{
Query: "q",
AgentID: "a",
TenantID: "t",
DocTypes: []string{"context"},
MaxResults: 10,
})
if err != nil {
t.Fatalf("search: %v", err)
}
for _, r := range res {
if r.Source == "kg" {
t.Fatalf("KG result leaked when types=context: %+v", r)
}
if r.Source == "episodic" {
t.Fatalf("episodic result leaked when types=context: %+v", r)
}
}
if kg.calls != 0 {
t.Fatalf("kg store should not be called when DocTypes=[context]; got %d calls", kg.calls)
}
if es.calls != 0 {
t.Fatalf("episodic store should not be called when DocTypes=[context]; got %d calls", es.calls)
}
}
// --- Test 1b: empty DocTypes still fans out to all sources. ---
func TestSearch_TypesEmptyIncludesAll(t *testing.T) {
vs := &fakeVaultStoreSearch{results: []store.VaultSearchResult{makeVaultResult("V1", "VaultDoc")}}
es := &fakeEpisodicStoreSearch{results: []store.EpisodicSearchResult{makeEpisodicResult("E1", "sess")}}
kg := &fakeKGStoreSearch{entities: []store.Entity{makeEntity("K1", "KGDoc")}}
svc := NewVaultSearchService(vs, es, kg)
res, err := svc.Search(context.Background(), UnifiedSearchOptions{
Query: "q",
AgentID: "a",
TenantID: "t",
MaxResults: 10,
})
if err != nil {
t.Fatalf("search: %v", err)
}
sources := map[string]bool{}
for _, r := range res {
sources[r.Source] = true
}
if !sources["vault"] || !sources["episodic"] || !sources["kg"] {
t.Fatalf("expected all 3 sources, got: %v", sources)
}
if kg.calls != 1 || es.calls != 1 || vs.calls != 1 {
t.Fatalf("expected 1 call each; got vault=%d ep=%d kg=%d", vs.calls, es.calls, kg.calls)
}
}
// --- Test 1c: DocTypes includes "kg" → kg fan-out runs. ---
func TestSearch_TypesIncludesKG(t *testing.T) {
vs := &fakeVaultStoreSearch{}
kg := &fakeKGStoreSearch{entities: []store.Entity{makeEntity("K1", "KGDoc")}}
svc := NewVaultSearchService(vs, nil, kg)
res, err := svc.Search(context.Background(), UnifiedSearchOptions{
Query: "q",
AgentID: "a",
TenantID: "t",
DocTypes: []string{"kg"},
MaxResults: 10,
})
if err != nil {
t.Fatalf("search: %v", err)
}
if kg.calls != 1 {
t.Fatalf("kg should run when DocTypes=[kg]; calls=%d", kg.calls)
}
found := false
for _, r := range res {
if r.Source == "kg" {
found = true
}
}
if !found {
t.Fatalf("expected kg result, got: %+v", res)
}
}