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>

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.

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))
 	}
 }

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

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)
+	}
+}

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))

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)
+	}
+}

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 {

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)
+	}
+}

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)
+	}
+}

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 {