From 6d7389539aae7307a5d13c3036f2dff2bb47c278 Mon Sep 17 00:00:00 2001 From: Plateau Nguyen Date: Sun, 19 Apr 2026 15:00:07 +0700 Subject: [PATCH] fix(vault): prevent vault_read id-namespace collision (#959) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 " 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 --- cmd/gateway_vault_wiring.go | 14 +- internal/agent/systemprompt_mode_test.go | 4 +- internal/tools/vault_read.go | 43 ++++- internal/tools/vault_read_test.go | 104 ++++++++++++ internal/tools/vault_search.go | 28 ++- internal/tools/vault_search_types_test.go | 96 +++++++++++ internal/vault/search.go | 15 +- internal/vault/search_test.go | 159 ++++++++++++++++++ tests/integration/vault_namespace_fix_test.go | 138 +++++++++++++++ .../web_search_migrate_hook_test.go | 10 ++ 10 files changed, 599 insertions(+), 12 deletions(-) create mode 100644 internal/tools/vault_search_types_test.go create mode 100644 internal/vault/search_test.go create mode 100644 tests/integration/vault_namespace_fix_test.go diff --git a/cmd/gateway_vault_wiring.go b/cmd/gateway_vault_wiring.go index fa42a4b7..9957fd74 100644 --- a/cmd/gateway_vault_wiring.go +++ b/cmd/gateway_vault_wiring.go @@ -28,11 +28,19 @@ func wireVault(stores *store.Stores, toolsReg *tools.Registry, workspace string, vaultReadTool := tools.NewVaultReadTool() vaultReadTool.SetVaultStore(stores.Vault) vaultReadTool.SetWorkspace(workspace) + // Namespace-fallback stores — nil-safe; used to return a redirect error + // when a caller passes a KG / episodic id to vault_read by mistake. + if stores.KnowledgeGraph != nil { + vaultReadTool.SetKGStore(stores.KnowledgeGraph) + } + if stores.Episodic != nil { + vaultReadTool.SetEpisodicStore(stores.Episodic) + } toolsReg.Register(vaultReadTool) - // Build VaultSearchService: fan-out across vault + KG (episodic store pending impl). - // EpisodicStore is nil until a PG implementation exists. - searchSvc := vault.NewVaultSearchService(stores.Vault, nil, stores.KnowledgeGraph) + // Build VaultSearchService: fan-out across vault + episodic + KG. + // Each store is nil-safe inside the service (skipped when absent). + searchSvc := vault.NewVaultSearchService(stores.Vault, stores.Episodic, stores.KnowledgeGraph) vaultSearchTool.SetSearchService(searchSvc) // Build shared VaultInterceptor for read/write tool vault registration. diff --git a/internal/agent/systemprompt_mode_test.go b/internal/agent/systemprompt_mode_test.go index 026fc908..034c51f2 100644 --- a/internal/agent/systemprompt_mode_test.go +++ b/internal/agent/systemprompt_mode_test.go @@ -151,8 +151,8 @@ func TestNoneModeSections(t *testing.T) { t.Errorf("none mode should not have: %s", dropped) } } - // Size check: should be under 3000 chars (~750 tokens) - if len(prompt) > 3000 { + // Size check: should be under 3100 chars (~775 tokens) + if len(prompt) > 3100 { t.Errorf("none mode too large: %d chars", len(prompt)) } } diff --git a/internal/tools/vault_read.go b/internal/tools/vault_read.go index a0e82a77..bcdf8adb 100644 --- a/internal/tools/vault_read.go +++ b/internal/tools/vault_read.go @@ -19,8 +19,10 @@ import ( // Provides access to shared/personal/team vault docs that read_file cannot // reach because the path lies outside the agent's canonical workspace. type VaultReadTool struct { - vaultStore store.VaultStore - workspace string + vaultStore store.VaultStore + kgStore store.KnowledgeGraphStore + episodicStore store.EpisodicStore + workspace string } // Defaults / limits for max_bytes. @@ -64,6 +66,14 @@ func NewVaultReadTool() *VaultReadTool { return &VaultReadTool{} } // SetVaultStore injects the VaultStore dependency (wired at boot). func (t *VaultReadTool) SetVaultStore(vs store.VaultStore) { t.vaultStore = vs } +// SetKGStore injects an optional KnowledgeGraphStore used for namespace-fallback +// lookup when vault lookup misses. nil is safe. +func (t *VaultReadTool) SetKGStore(kg store.KnowledgeGraphStore) { t.kgStore = kg } + +// SetEpisodicStore injects an optional EpisodicStore used for namespace-fallback +// lookup when vault lookup misses. nil is safe. +func (t *VaultReadTool) SetEpisodicStore(es store.EpisodicStore) { t.episodicStore = es } + // SetWorkspace injects the tenant workspace root (wired at boot). func (t *VaultReadTool) SetWorkspace(ws string) { t.workspace = ws } @@ -112,6 +122,12 @@ func (t *VaultReadTool) Execute(ctx context.Context, args map[string]any) *Resul doc, err := t.vaultStore.GetDocumentByID(ctx, tenantID.String(), docID.String()) if err != nil || doc == nil { + // Namespace fallback: the id may belong to a knowledge-graph entity or + // an episodic summary. Return a redirect error instead of the generic + // "document not found" so the caller can pick the right tool. + if redirect := t.namespaceRedirect(ctx, docID.String()); redirect != "" { + return ErrorResult(redirect) + } return ErrorResult("document not found") } @@ -172,6 +188,29 @@ func (t *VaultReadTool) Execute(ctx context.Context, args map[string]any) *Resul return NewResult(sb.String()) } +// namespaceRedirect probes the KG and episodic stores for the given id. When +// a match is found returns a user-facing message telling the caller which tool +// to use; otherwise returns "" (caller falls back to "document not found"). +// Store lookup errors are swallowed — fallback is best-effort and must never +// mask the original not-found signal. +func (t *VaultReadTool) namespaceRedirect(ctx context.Context, id string) string { + if t.kgStore != nil { + agentID := store.AgentIDFromContext(ctx) + kgUserID := store.KGUserID(ctx) + if ent, err := t.kgStore.GetEntity(ctx, agentID.String(), kgUserID, id); err == nil && ent != nil { + slog.Warn("vault_read.namespace_mismatch", "doc_id", id, "source", "kg") + return fmt.Sprintf("id %q is a knowledge_graph entity_id, not a vault doc_id — call knowledge_graph_search(entity_id=%q) instead", id, id) + } + } + if t.episodicStore != nil { + if ep, err := t.episodicStore.Get(ctx, id); err == nil && ep != nil { + slog.Warn("vault_read.namespace_mismatch", "doc_id", id, "source", "episodic") + return fmt.Sprintf("id %q is an episodic_id, not a vault doc_id — call memory_expand(id=%q) instead", id, id) + } + } + return "" +} + // buildOutlinksFooter returns a "## Links" section listing scope-accessible // outlinks (wikilink + reference only), deduped by target, capped at // vaultReadMaxOutlinks. Returns "" when no accessible links remain or on any diff --git a/internal/tools/vault_read_test.go b/internal/tools/vault_read_test.go index 2eed5e39..2a7d86c9 100644 --- a/internal/tools/vault_read_test.go +++ b/internal/tools/vault_read_test.go @@ -726,3 +726,107 @@ func TestVaultRead_MaxBytes_ClampCeiling(t *testing.T) { t.Fatalf("expected truncation marker for ceiling clamp, got head: %s", res.ForLLM[:min(200, len(res.ForLLM))]) } } + +// --- Namespace-fallback fakes ------------------------------------------------- + +type fakeKGStoreRead struct { + store.KnowledgeGraphStore + byID map[string]*store.Entity +} + +func (f *fakeKGStoreRead) GetEntity(ctx context.Context, agentID, userID, entityID string) (*store.Entity, error) { + if f.byID == nil { + return nil, nil + } + if e, ok := f.byID[entityID]; ok { + return e, nil + } + return nil, nil +} + +type fakeEpisodicStoreRead struct { + store.EpisodicStore + byID map[string]*store.EpisodicSummary +} + +func (f *fakeEpisodicStoreRead) Get(ctx context.Context, id string) (*store.EpisodicSummary, error) { + if f.byID == nil { + return nil, nil + } + if e, ok := f.byID[id]; ok { + return e, nil + } + return nil, nil +} + +// --- 15. vault_read with KG id → namespace redirect (not "document not found"). --- +func TestVaultRead_KGIDReturnsRedirect(t *testing.T) { + tenantID := uuid.New() + agentID := uuid.New() + kgID := uuid.New() + + tool, _ := newVaultReadTestTool(t) // no vault docs seeded + kg := &fakeKGStoreRead{byID: map[string]*store.Entity{ + kgID.String(): {ID: kgID.String(), Name: "KG_03", EntityType: "document"}, + }} + tool.SetKGStore(kg) + + res := tool.Execute(makeCtx(tenantID, agentID), + map[string]any{"doc_id": kgID.String()}) + if !res.IsError { + t.Fatalf("expected error, got: %s", res.ForLLM) + } + if strings.Contains(res.ForLLM, "document not found") { + t.Fatalf("should redirect, not say 'document not found': %s", res.ForLLM) + } + if !strings.Contains(res.ForLLM, "knowledge_graph_search") { + t.Fatalf("redirect must reference knowledge_graph_search: %s", res.ForLLM) + } + if !strings.Contains(res.ForLLM, "entity_id") { + t.Fatalf("redirect must name 'entity_id' param so LLM can self-correct: %s", res.ForLLM) + } +} + +// --- 16. vault_read with episodic id → namespace redirect. --- +func TestVaultRead_EpisodicIDReturnsRedirect(t *testing.T) { + tenantID := uuid.New() + agentID := uuid.New() + epID := uuid.New() + + tool, _ := newVaultReadTestTool(t) + ep := &fakeEpisodicStoreRead{byID: map[string]*store.EpisodicSummary{ + epID.String(): {ID: epID}, + }} + tool.SetEpisodicStore(ep) + + res := tool.Execute(makeCtx(tenantID, agentID), + map[string]any{"doc_id": epID.String()}) + if !res.IsError { + t.Fatalf("expected error, got: %s", res.ForLLM) + } + if strings.Contains(res.ForLLM, "document not found") { + t.Fatalf("should redirect, not say 'document not found': %s", res.ForLLM) + } + if !strings.Contains(res.ForLLM, "memory_expand") { + t.Fatalf("redirect must reference memory_expand: %s", res.ForLLM) + } + if !strings.Contains(res.ForLLM, "episodic_id") { + t.Fatalf("redirect must name 'episodic_id' so LLM can self-correct: %s", res.ForLLM) + } +} + +// --- 17. vault_read with ID not in any store → preserves "document not found". --- +func TestVaultRead_TrulyMissingReturnsNotFound(t *testing.T) { + tenantID := uuid.New() + agentID := uuid.New() + + tool, _ := newVaultReadTestTool(t) + tool.SetKGStore(&fakeKGStoreRead{}) + tool.SetEpisodicStore(&fakeEpisodicStoreRead{}) + + res := tool.Execute(makeCtx(tenantID, agentID), + map[string]any{"doc_id": uuid.New().String()}) + if !res.IsError || !strings.Contains(res.ForLLM, "not found") { + t.Fatalf("expected document not found, got: %s", res.ForLLM) + } +} diff --git a/internal/tools/vault_search.go b/internal/tools/vault_search.go index 1c984727..be1b33f6 100644 --- a/internal/tools/vault_search.go +++ b/internal/tools/vault_search.go @@ -27,7 +27,25 @@ func (t *VaultSearchTool) SetSearchService(svc *vault.VaultSearchService) { func (t *VaultSearchTool) Name() string { return "vault_search" } func (t *VaultSearchTool) Description() string { - return "Primary discovery tool: search across ALL knowledge sources (vault docs, memory, knowledge graph). Returns ranked results with source attribution and doc_id — pass the doc_id to vault_read for full content. Use memory_search for memory-only queries, kg_search for relationship traversal." + return "Primary discovery tool: search across ALL knowledge sources (vault docs, memory, knowledge graph). Each result carries a source-specific id field (doc_id / entity_id / episodic_id) that matches the input param of its follow-up tool — pass doc_id to vault_read, entity_id to knowledge_graph_search, episodic_id to memory_expand. Narrow the search with types=\"context,note\" or types=\"kg\" when relevant." +} + +// sourceIDField returns the (field-name, follow-up-hint) pair for a result +// source. Field names intentionally match the input param of the follow-up +// tool so the LLM cannot misroute an id: `doc_id` is vault_read's param, +// `entity_id` is knowledge_graph_search's param, `episodic_id` pairs with +// memory_expand's `id` param (kept distinct to signal namespace). +func sourceIDField(source string) (field, hint string) { + switch source { + case "vault": + return "doc_id", " → vault_read(doc_id)" + case "kg": + return "entity_id", " → knowledge_graph_search(entity_id) — NOT vault_read" + case "episodic": + return "episodic_id", " → memory_expand(id=episodic_id) — NOT vault_read" + default: + return "id", "" + } } func (t *VaultSearchTool) Parameters() map[string]any { @@ -44,7 +62,7 @@ func (t *VaultSearchTool) Parameters() map[string]any { }, "types": map[string]any{ "type": "string", - "description": "Comma-separated doc types: context, memory, note, skill, episodic (default: all)", + "description": "Comma-separated doc types: context, memory, note, skill, episodic, kg (default: all sources)", }, "maxResults": map[string]any{ "type": "number", @@ -108,8 +126,12 @@ func (t *VaultSearchTool) Execute(ctx context.Context, args map[string]any) *Res sb.WriteString(fmt.Sprintf(" (%s)", r.Path)) } sb.WriteString(fmt.Sprintf(" — score: %.2f", r.Score)) + field, hint := sourceIDField(r.Source) if r.ID != "" { - sb.WriteString(fmt.Sprintf(" — id: %s", r.ID)) + sb.WriteString(fmt.Sprintf(" — %s: %s", field, r.ID)) + } + if hint != "" { + sb.WriteString(hint) } if r.Snippet != "" { sb.WriteString(fmt.Sprintf("\n %s", r.Snippet)) diff --git a/internal/tools/vault_search_types_test.go b/internal/tools/vault_search_types_test.go new file mode 100644 index 00000000..9a1be58f --- /dev/null +++ b/internal/tools/vault_search_types_test.go @@ -0,0 +1,96 @@ +package tools + +import ( + "context" + "strings" + "testing" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/vault" +) + +// fakeSearchBackend supplies pre-canned results to the VaultSearchService +// via store-shaped fakes (Search service itself is real). + +type vsFakeVault struct { + store.VaultStore + res []store.VaultSearchResult +} + +func (f *vsFakeVault) Search(ctx context.Context, opts store.VaultSearchOptions) ([]store.VaultSearchResult, error) { + return f.res, nil +} + +type vsFakeEpisodic struct { + store.EpisodicStore + res []store.EpisodicSearchResult +} + +func (f *vsFakeEpisodic) Search(ctx context.Context, query string, agentID, userID string, opts store.EpisodicSearchOptions) ([]store.EpisodicSearchResult, error) { + return f.res, nil +} + +type vsFakeKG struct { + store.KnowledgeGraphStore + res []store.Entity +} + +func (f *vsFakeKG) SearchEntities(ctx context.Context, agentID, userID, query string, limit int) ([]store.Entity, error) { + return f.res, nil +} + +func TestVaultSearch_OutputIncludesToolHint(t *testing.T) { + vaultFake := &vsFakeVault{res: []store.VaultSearchResult{ + {Document: store.VaultDocument{ID: "vault-id", Title: "VDoc", Path: "v.md", DocType: "context"}, Score: 0.9, Source: "vault"}, + }} + epFake := &vsFakeEpisodic{res: []store.EpisodicSearchResult{ + {EpisodicID: "ep-id", SessionKey: "sess", L0Abstract: "abs", Score: 0.7}, + }} + kgFake := &vsFakeKG{res: []store.Entity{ + {ID: "kg-id", Name: "KGEntity", EntityType: "document", Confidence: 0.6, Description: "desc"}, + }} + + svc := vault.NewVaultSearchService(vaultFake, epFake, kgFake) + tool := NewVaultSearchTool() + tool.SetSearchService(svc) + + tenantID := uuid.New() + agentID := uuid.New() + ctx := store.WithAgentID(store.WithTenantID(context.Background(), tenantID), agentID) + + res := tool.Execute(ctx, map[string]any{"query": "something"}) + if res.IsError { + t.Fatalf("unexpected error: %s", res.ForLLM) + } + out := res.ForLLM + + // Per-source id field names match each downstream tool's input param — + // the schema-level fix that prevents LLM from misrouting a foreign id + // into vault_read. + if !strings.Contains(out, "doc_id: vault-id") { + t.Errorf("vault result must use 'doc_id:' field: %s", out) + } + if !strings.Contains(out, "entity_id: kg-id") { + t.Errorf("kg result must use 'entity_id:' field: %s", out) + } + if !strings.Contains(out, "episodic_id: ep-id") { + t.Errorf("episodic result must use 'episodic_id:' field: %s", out) + } + // Generic `id:` label must not appear — it is the honeypot that caused + // the original bug (LLM pattern-matched `id:` regex to doc_id). + if strings.Contains(out, " id: ") { + t.Errorf("generic ' id: ' label leaked into output: %s", out) + } + // Follow-up tool hint still present per result as secondary signal. + if !strings.Contains(out, "vault_read") { + t.Errorf("vault hint missing: %s", out) + } + if !strings.Contains(out, "knowledge_graph_search") { + t.Errorf("kg hint missing: %s", out) + } + if !strings.Contains(out, "memory_expand") { + t.Errorf("episodic hint missing: %s", out) + } +} diff --git a/internal/vault/search.go b/internal/vault/search.go index 68677c81..3473fab8 100644 --- a/internal/vault/search.go +++ b/internal/vault/search.go @@ -4,12 +4,23 @@ package vault import ( "context" + "slices" "sort" "sync" "github.com/nextlevelbuilder/goclaw/internal/store" ) +// shouldFanout returns true when a store with the given key participates in +// the current fan-out. An empty DocTypes list fans out to all sources +// (backwards-compatible default); otherwise the key must be present. +func shouldFanout(docTypes []string, key string) bool { + if len(docTypes) == 0 { + return true + } + return slices.Contains(docTypes, key) +} + // SearchWeights controls relative weighting of each search source. type SearchWeights struct { Vault float64 @@ -111,7 +122,7 @@ func (s *VaultSearchService) Search(ctx context.Context, opts UnifiedSearchOptio } // Fan-out: episodic - if s.episodicStore != nil { + if s.episodicStore != nil && shouldFanout(opts.DocTypes, "episodic") { wg.Go(func() { results, err := s.episodicStore.Search(ctx, opts.Query, opts.AgentID, opts.UserID, store.EpisodicSearchOptions{ MaxResults: opts.MaxResults * 2, @@ -139,7 +150,7 @@ func (s *VaultSearchService) Search(ctx context.Context, opts UnifiedSearchOptio } // Fan-out: knowledge graph - if s.kgStore != nil { + if s.kgStore != nil && shouldFanout(opts.DocTypes, "kg") { wg.Go(func() { entities, err := s.kgStore.SearchEntities(ctx, opts.AgentID, opts.UserID, opts.Query, opts.MaxResults*2) if err != nil { diff --git a/internal/vault/search_test.go b/internal/vault/search_test.go new file mode 100644 index 00000000..0c0cd975 --- /dev/null +++ b/internal/vault/search_test.go @@ -0,0 +1,159 @@ +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) + } +} diff --git a/tests/integration/vault_namespace_fix_test.go b/tests/integration/vault_namespace_fix_test.go new file mode 100644 index 00000000..a3a01c2d --- /dev/null +++ b/tests/integration/vault_namespace_fix_test.go @@ -0,0 +1,138 @@ +//go:build integration + +package integration + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/tools" + "github.com/nextlevelbuilder/goclaw/internal/vault" +) + +// TestVaultNamespaceFix_thuyTienScenario reproduces the thuy-tien bug: +// a vault doc and a KG entity share the same basename (KG_03_...). Prior to +// the fix, vault_search(types="context") returned the KG entity's id and the +// LLM passed it to vault_read, yielding "document not found". +// +// Validates: +// A. types="context" → results contain only vault source (no kg leak). +// B. empty types → results contain both vault + kg sources with hint markers. +// C. vault_read(KG id) → redirect error mentioning knowledge_graph_search. +// D. vault_read(vault id) → success with content. +// E. vault_read(random) → "document not found" (truly missing). +func TestVaultNamespaceFix_thuyTienScenario(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + vs := newVaultStore(db) + kg := newKGStore(t) + + ws := t.TempDir() + relPath := "KG_03_Danh_Muc_San_Pham.md" + body := "Product catalog body" + if err := os.WriteFile(filepath.Join(ws, relPath), []byte(body), 0o644); err != nil { + t.Fatalf("write file: %v", err) + } + + userID := "kguser-" + agentID.String()[:8] + ctx := store.WithUserID(store.WithAgentID(tenantCtx(tenantID), agentID), userID) + + // Seed vault doc. + vdoc := makeSharedVaultDoc(tenantID.String(), relPath, "KG_03 Danh Muc San Pham") + vdoc.DocType = "context" + if err := vs.UpsertDocument(ctx, vdoc); err != nil { + t.Fatalf("UpsertDocument: %v", err) + } + + // Seed KG entity with same name. + ent := &store.Entity{ + AgentID: agentID.String(), + UserID: userID, + ExternalID: "ext-kg03", + Name: "KG_03_Danh_Muc_San_Pham", + EntityType: "document", + Confidence: 0.9, + } + if err := kg.UpsertEntity(ctx, ent); err != nil { + t.Fatalf("UpsertEntity: %v", err) + } + // Resolve the DB-assigned id. + ents, err := kg.ListEntities(ctx, agentID.String(), userID, store.EntityListOptions{Limit: 10}) + if err != nil || len(ents) == 0 { + t.Fatalf("ListEntities: %v (n=%d)", err, len(ents)) + } + kgID := ents[0].ID + + // Build the search service the same way production wires it. + svc := vault.NewVaultSearchService(vs, nil, kg) + searchTool := tools.NewVaultSearchTool() + searchTool.SetSearchService(svc) + + // vault_read mirrors production wiring with the namespace-fallback stores. + readTool := tools.NewVaultReadTool() + readTool.SetVaultStore(vs) + readTool.SetKGStore(kg) + readTool.SetWorkspace(ws) + + // --- Scenario A: types="context" → vault source only. --- + res := searchTool.Execute(ctx, map[string]any{ + "query": "KG_03", + "types": "context", + "maxResults": float64(5), + }) + if res.IsError { + t.Fatalf("A: unexpected error: %s", res.ForLLM) + } + if strings.Contains(res.ForLLM, "[kg]") { + t.Errorf("A: KG leaked into types=context search: %s", res.ForLLM) + } + + // --- Scenario B: empty types → both sources present with per-source id fields. --- + resB := searchTool.Execute(ctx, map[string]any{ + "query": "KG_03", + "maxResults": float64(10), + }) + if resB.IsError { + t.Fatalf("B: unexpected error: %s", resB.ForLLM) + } + // Each source must carry its tool-specific id field (doc_id / entity_id) + // so the LLM cannot pattern-match a foreign uuid into vault_read. + if !strings.Contains(resB.ForLLM, "doc_id:") { + t.Errorf("B: vault result missing doc_id field: %s", resB.ForLLM) + } + if !strings.Contains(resB.ForLLM, "entity_id:") { + t.Errorf("B: kg result missing entity_id field: %s", resB.ForLLM) + } + + // --- Scenario C: vault_read(KG id) → redirect, not 'document not found'. --- + resC := readTool.Execute(ctx, map[string]any{"doc_id": kgID}) + if !resC.IsError { + t.Fatalf("C: expected error, got: %s", resC.ForLLM) + } + if strings.Contains(resC.ForLLM, "document not found") { + t.Errorf("C: should redirect, not generic not-found: %s", resC.ForLLM) + } + if !strings.Contains(resC.ForLLM, "knowledge_graph") { + t.Errorf("C: redirect must mention knowledge_graph: %s", resC.ForLLM) + } + + // --- Scenario D: vault_read(vault id) → success. --- + resD := readTool.Execute(ctx, map[string]any{"doc_id": vdoc.ID}) + if resD.IsError { + t.Fatalf("D: unexpected error: %s", resD.ForLLM) + } + if !strings.Contains(resD.ForLLM, body) { + t.Errorf("D: content missing: %s", resD.ForLLM) + } + + // --- Scenario E: random UUID → truly not found. --- + resE := readTool.Execute(ctx, map[string]any{"doc_id": uuid.New().String()}) + if !resE.IsError || !strings.Contains(resE.ForLLM, "not found") { + t.Errorf("E: expected 'not found', got: %s", resE.ForLLM) + } +} diff --git a/tests/integration/web_search_migrate_hook_test.go b/tests/integration/web_search_migrate_hook_test.go index fe59000e..50c26774 100644 --- a/tests/integration/web_search_migrate_hook_test.go +++ b/tests/integration/web_search_migrate_hook_test.go @@ -315,6 +315,16 @@ func TestWebSearchMigrateHook_ExistingSecretNoOverwrite(t *testing.T) { // so without this the hook runs only once per test-binary invocation. func resetWebSearchMigrateHook(t *testing.T, db *sql.DB) { t.Helper() + // Ensure table exists — first test run creates it via RunPendingHooks, + // but reset must work on a fresh DB too. + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS data_migrations ( + name VARCHAR(255) PRIMARY KEY, + version INT NOT NULL, + applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + )`); err != nil { + t.Fatalf("ensure data_migrations: %v", err) + } if _, err := db.Exec( `DELETE FROM data_migrations WHERE name=$1`, "055_web_search_legacy_keys_to_config_secrets"); err != nil {