feat(tools): append scope-filtered outlinks footer to vault_read

List up to 20 wikilink/reference outlinks as "## Links" after the doc
body, deduped by target, scope-filtered via existing allowed() matrix.
Non-fatal — store errors log and skip the footer.
This commit is contained in:
viettranx committed 2026-04-18 15:02:34 +07:00
1 parent e033b357cc
commit 001d9e23e3
3 files changed
+508 -4

No files matched your search

+99
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"io"
"log/slog"
"os"
"path/filepath"
"strings"
@@ -27,8 +28,16 @@ const (
vaultReadDefaultMaxBytes = 500_000
vaultReadCeilingMaxBytes = 1_048_576
vaultReadUTF8SniffBytes = 8192
vaultReadMaxOutlinks = 20
)
// outlinkAllowedTypes restricts the footer to knowledge-graph links; operational
// link types (task_attachment, delegation_attachment, ...) are excluded.
var outlinkAllowedTypes = map[string]struct{}{
"wikilink": {},
"reference": {},
}
// nonTextExtensions are file extensions rejected by the text-only gate.
// Lower-case with leading dot. Keep in sync with phase 01 spec.
var nonTextExtensions = map[string]struct{}{
@@ -159,9 +168,99 @@ func (t *VaultReadTool) Execute(ctx context.Context, args map[string]any) *Resul
if truncated {
fmt.Fprintf(&sb, "\n\n…[truncated, content exceeds %d bytes]", maxBytes)
}
sb.WriteString(t.buildOutlinksFooter(ctx, tenantID.String(), doc.ID))
return NewResult(sb.String())
}
// 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
// store error — footer is additive and must not break the read.
func (t *VaultReadTool) buildOutlinksFooter(ctx context.Context, tenantID, docID string) string {
links, err := t.vaultStore.GetOutLinks(ctx, tenantID, docID)
if err != nil {
slog.Warn("vault_read.outlinks.get_failed", "doc_id", docID, "error", err)
return ""
}
if len(links) == 0 {
return ""
}
// Pass 1: type filter + self-link drop + dedup by ToDocID.
seen := make(map[string]struct{}, len(links))
filtered := make([]store.VaultLink, 0, len(links))
for _, l := range links {
if _, ok := outlinkAllowedTypes[l.LinkType]; !ok {
continue
}
if l.ToDocID == docID {
continue
}
if _, ok := seen[l.ToDocID]; ok {
continue
}
seen[l.ToDocID] = struct{}{}
filtered = append(filtered, l)
}
if len(filtered) == 0 {
return ""
}
// Batch fetch targets.
ids := make([]string, 0, len(filtered))
for _, l := range filtered {
ids = append(ids, l.ToDocID)
}
targets, err := t.vaultStore.GetDocumentsByIDs(ctx, tenantID, ids)
if err != nil {
slog.Warn("vault_read.outlinks.targets_failed", "doc_id", docID, "error", err)
return ""
}
byID := make(map[string]store.VaultDocument, len(targets))
for _, d := range targets {
byID[d.ID] = d
}
// Pass 2: missing-target + scope filter, cap.
type keptLink struct {
doc store.VaultDocument
linkType string
}
kept := make([]keptLink, 0, vaultReadMaxOutlinks)
var overflow int
for _, l := range filtered {
target, ok := byID[l.ToDocID]
if !ok {
continue
}
if !t.allowed(ctx, &target) {
continue
}
if len(kept) >= vaultReadMaxOutlinks {
overflow++
continue
}
kept = append(kept, keptLink{doc: target, linkType: l.LinkType})
}
if len(kept) == 0 {
return ""
}
var sb strings.Builder
sb.WriteString("\n\n## Links\n")
for _, k := range kept {
title := k.doc.Title
if title == "" {
title = k.doc.Path
}
fmt.Fprintf(&sb, "- %s — id: %s (%s)\n", title, k.doc.ID, k.linkType)
}
if overflow > 0 {
fmt.Fprintf(&sb, "…[%d more links omitted]\n", overflow)
}
return sb.String()
}
// allowed enforces the scope matrix:
// - shared: allow.
// - personal: allow iff doc.AgentID == agentID from ctx.
+330 -4
View File
@@ -3,6 +3,7 @@ package tools
import (
"bytes"
"context"
"fmt"
"os"
"path/filepath"
"strings"
@@ -14,12 +15,14 @@ import (
)
// fakeVaultStore embeds store.VaultStore (nil) so the struct satisfies the
// interface at compile time; only GetDocumentByID is implemented for tests.
// Any call to an unimplemented method would nil-pointer panic — which is fine
// because vault_read.Execute only hits GetDocumentByID.
// interface at compile time. Methods used by vault_read.Execute are
// implemented; others would nil-panic if called.
type fakeVaultStore struct {
store.VaultStore
byID map[string]*store.VaultDocument // key: tenantID + ":" + docID
byID map[string]*store.VaultDocument // key: tenantID + ":" + docID
outlinks map[string][]store.VaultLink // key: tenantID + ":" + docID
outLinkErr error // injected error for GetOutLinks
targetsErr error // injected error for GetDocumentsByIDs
}
func (f *fakeVaultStore) GetDocumentByID(ctx context.Context, tenantID, id string) (*store.VaultDocument, error) {
@@ -33,6 +36,29 @@ func (f *fakeVaultStore) GetDocumentByID(ctx context.Context, tenantID, id strin
return doc, nil
}
func (f *fakeVaultStore) GetOutLinks(ctx context.Context, tenantID, docID string) ([]store.VaultLink, error) {
if f.outLinkErr != nil {
return nil, f.outLinkErr
}
if f.outlinks == nil {
return nil, nil
}
return f.outlinks[tenantID+":"+docID], nil
}
func (f *fakeVaultStore) GetDocumentsByIDs(ctx context.Context, tenantID string, docIDs []string) ([]store.VaultDocument, error) {
if f.targetsErr != nil {
return nil, f.targetsErr
}
out := make([]store.VaultDocument, 0, len(docIDs))
for _, id := range docIDs {
if d, ok := f.byID[tenantID+":"+id]; ok {
out = append(out, *d)
}
}
return out, nil
}
// newVaultReadTestTool builds a VaultReadTool with a temp workspace and a
// fake store pre-seeded with docs.
func newVaultReadTestTool(t *testing.T, docs ...*store.VaultDocument) (*VaultReadTool, string) {
@@ -375,6 +401,306 @@ func TestVaultRead_SymlinkEscape_Denied(t *testing.T) {
}
}
// --- Outlinks footer tests -------------------------------------------------
// sharedDoc returns a minimal shared-scope vault doc.
func sharedDoc(tenantID uuid.UUID, title, path string) *store.VaultDocument {
return &store.VaultDocument{
ID: uuid.New().String(), TenantID: tenantID.String(),
Scope: "shared", Path: path, Title: title, DocType: "note",
}
}
// personalDoc returns a personal-scope doc owned by agentID.
func personalDoc(tenantID, agentID uuid.UUID, title, path string) *store.VaultDocument {
aid := agentID.String()
return &store.VaultDocument{
ID: uuid.New().String(), TenantID: tenantID.String(),
AgentID: &aid, Scope: "personal", Path: path, Title: title, DocType: "note",
}
}
// teamDoc returns a team-scope doc bound to teamID.
func teamDoc(tenantID uuid.UUID, teamID, title, path string) *store.VaultDocument {
tid := teamID
return &store.VaultDocument{
ID: uuid.New().String(), TenantID: tenantID.String(),
TeamID: &tid, Scope: "team", Path: path, Title: title, DocType: "note",
}
}
// seedWithLinks extends the fake store's outlinks map.
func seedLinks(tool *VaultReadTool, tenantID, fromID string, links []store.VaultLink) {
f := tool.vaultStore.(*fakeVaultStore)
if f.outlinks == nil {
f.outlinks = make(map[string][]store.VaultLink)
}
f.outlinks[tenantID+":"+fromID] = links
}
// --- Case 1: outlinks present, all in scope → listed in order. ---
func TestVaultReadOutlinks_AllInScope_Listed(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tA := sharedDoc(tenantID, "Target A", "a.md")
tB := sharedDoc(tenantID, "Target B", "b.md")
tool, ws := newVaultReadTestTool(t, src, tA, tB)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: tA.ID, LinkType: "wikilink"},
{ToDocID: tB.ID, LinkType: "reference"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("missing Links heading: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Target A — id: "+tA.ID+" (wikilink)") {
t.Fatalf("target A line missing/wrong: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Target B — id: "+tB.ID+" (reference)") {
t.Fatalf("target B line missing/wrong: %s", res.ForLLM)
}
// Order preserved (A before B).
if strings.Index(res.ForLLM, "Target A") > strings.Index(res.ForLLM, "Target B") {
t.Fatalf("order not preserved: %s", res.ForLLM)
}
}
// --- Case 2: outlink to personal doc owned by another agent → dropped. ---
func TestVaultReadOutlinks_PersonalOtherAgent_Dropped(t *testing.T) {
tenantID := uuid.New()
agentSelf := uuid.New()
agentOther := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
other := personalDoc(tenantID, agentOther, "Other Memo", "other.md")
tool, ws := newVaultReadTestTool(t, src, other)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: other.ID, LinkType: "wikilink"},
})
res := tool.Execute(makeCtx(tenantID, agentSelf), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "Other Memo") {
t.Fatalf("leaked other-agent personal doc: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("empty Links heading should not appear: %s", res.ForLLM)
}
}
// --- Case 3: outlink to team doc from different team → dropped. ---
func TestVaultReadOutlinks_TeamOtherTeam_Dropped(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
teamSelf := uuid.New().String()
teamOther := uuid.New().String()
src := sharedDoc(tenantID, "Source", "src.md")
target := teamDoc(tenantID, teamOther, "Other Team Doc", "ot.md")
tool, ws := newVaultReadTestTool(t, src, target)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: target.ID, LinkType: "wikilink"},
})
res := tool.Execute(makeCtxWithTeam(tenantID, agentID, teamSelf),
map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "Other Team Doc") {
t.Fatalf("leaked other-team doc: %s", res.ForLLM)
}
}
// --- Case 4: outlink to shared doc → always included. ---
func TestVaultReadOutlinks_Shared_Included(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tgt := sharedDoc(tenantID, "Shared Target", "st.md")
tool, ws := newVaultReadTestTool(t, src, tgt)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: tgt.ID, LinkType: "wikilink"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Shared Target") {
t.Fatalf("shared target missing: %s", res.ForLLM)
}
}
// --- Case 5: task_attachment link type → dropped by type filter. ---
func TestVaultReadOutlinks_TaskAttachment_Dropped(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tgt := sharedDoc(tenantID, "Task Target", "tk.md")
tool, ws := newVaultReadTestTool(t, src, tgt)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: tgt.ID, LinkType: "task_attachment"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("footer should be absent for task_attachment only: %s", res.ForLLM)
}
}
// --- Case 6: broken link (target doc deleted) → dropped silently. ---
func TestVaultReadOutlinks_BrokenLink_Dropped(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
ghostID := uuid.New().String()
tool, ws := newVaultReadTestTool(t, src) // ghost not seeded
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: ghostID, LinkType: "wikilink"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("footer should be absent when only broken links: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, ghostID) {
t.Fatalf("ghost id leaked: %s", res.ForLLM)
}
}
// --- Case 7: self-link → dropped. ---
func TestVaultReadOutlinks_SelfLink_Dropped(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tool, ws := newVaultReadTestTool(t, src)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: src.ID, LinkType: "wikilink"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("self-link should not produce footer: %s", res.ForLLM)
}
}
// --- Case 8: >20 valid links → capped with overflow marker. ---
func TestVaultReadOutlinks_Overflow_Capped(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
docs := []*store.VaultDocument{src}
links := []store.VaultLink{}
for i := range 25 {
d := sharedDoc(tenantID, fmt.Sprintf("T%02d", i), fmt.Sprintf("t%02d.md", i))
docs = append(docs, d)
links = append(links, store.VaultLink{ToDocID: d.ID, LinkType: "wikilink"})
}
tool, ws := newVaultReadTestTool(t, docs...)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, links)
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "…[5 more links omitted]") {
t.Fatalf("expected overflow marker for 25→20 cap: %s", res.ForLLM)
}
// Count rendered link lines (prefix "- T").
n := strings.Count(res.ForLLM, "\n- T")
if n != 20 {
t.Fatalf("expected 20 kept links, got %d", n)
}
}
// --- Case 9: zero valid links → no heading at all. ---
func TestVaultReadOutlinks_ZeroValid_NoHeading(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tool, ws := newVaultReadTestTool(t, src)
writeFile(t, ws, "src.md", "body")
// no links seeded.
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("should not emit empty heading: %s", res.ForLLM)
}
}
// --- Case 10: dedup same target via wikilink + reference → shown once. ---
func TestVaultReadOutlinks_Dedup_ByTarget(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tgt := sharedDoc(tenantID, "DupTarget", "dup.md")
tool, ws := newVaultReadTestTool(t, src, tgt)
writeFile(t, ws, "src.md", "body")
seedLinks(tool, tenantID.String(), src.ID, []store.VaultLink{
{ToDocID: tgt.ID, LinkType: "wikilink"},
{ToDocID: tgt.ID, LinkType: "reference"},
})
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if n := strings.Count(res.ForLLM, "DupTarget"); n != 1 {
t.Fatalf("expected 1 occurrence of DupTarget, got %d: %s", n, res.ForLLM)
}
// First link_type wins.
if !strings.Contains(res.ForLLM, "(wikilink)") {
t.Fatalf("expected first link_type wikilink to win: %s", res.ForLLM)
}
}
// --- Case: GetOutLinks error → read succeeds, footer empty. ---
func TestVaultReadOutlinks_StoreError_NoFooter(t *testing.T) {
tenantID := uuid.New()
agentID := uuid.New()
src := sharedDoc(tenantID, "Source", "src.md")
tool, ws := newVaultReadTestTool(t, src)
writeFile(t, ws, "src.md", "body")
tool.vaultStore.(*fakeVaultStore).outLinkErr = os.ErrInvalid
res := tool.Execute(makeCtx(tenantID, agentID), map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("store error must not fail read: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "body") {
t.Fatalf("body missing: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("footer must be absent on store error: %s", res.ForLLM)
}
}
// --- 14. max_bytes clamp to ceiling (1MB). ---
func TestVaultRead_MaxBytes_ClampCeiling(t *testing.T) {
tenantID := uuid.New()
+79
View File
@@ -176,3 +176,82 @@ func TestVaultRead_OversizeTruncation(t *testing.T) {
// Issue #948 regression guard: identifier in test name helps future greps.
_ = uuid.Nil
}
// TestVaultRead_OutlinksScopeMatrix verifies the inline ## Links footer honours
// the scope matrix end-to-end with a real VaultStore:
// - shared target → included
// - personal target (self) → included
// - personal target (other)→ dropped
// - team target (other) → dropped
//
// Uses wikilink type so the type-filter (wikilink+reference only) lets all
// test links through; the scope matrix is what we are exercising.
func TestVaultRead_OutlinksScopeMatrix(t *testing.T) {
db := testDB(t)
tenantID, agentSelf := seedTenantAgent(t, db)
_, agentOther := seedTenantAgent(t, db)
vs := newVaultStore(db)
ws := t.TempDir()
tid := tenantID.String()
ctxSelf := store.WithAgentID(tenantCtx(tenantID), agentSelf)
// Source file on disk + doc (shared so read is allowed).
srcRel := "src.md"
if err := os.WriteFile(filepath.Join(ws, srcRel), []byte("src body"), 0o644); err != nil {
t.Fatalf("write src: %v", err)
}
src := makeSharedVaultDoc(tid, srcRel, "Source")
if err := vs.UpsertDocument(ctxSelf, src); err != nil {
t.Fatalf("Upsert src: %v", err)
}
// Targets.
// Seed a real team belonging to a DIFFERENT agent so the "team-other"
// target satisfies FK and is scope-denied from agentSelf's RunContext.
otherTeamID, _ := seedTeam(t, db, tenantID, agentOther)
shared := makeSharedVaultDoc(tid, "shared.md", "SharedTgt")
personalSelf := makeVaultDoc(tid, agentSelf.String(), "self.md", "SelfTgt")
personalOther := makeVaultDoc(tid, agentOther.String(), "other.md", "PersonalOtherTgt")
teamOther := makeTeamVaultDoc(tid, otherTeamID.String(), "team.md", "TeamOtherTgt")
for _, d := range []*store.VaultDocument{shared, personalSelf, personalOther, teamOther} {
if err := vs.UpsertDocument(ctxSelf, d); err != nil {
t.Fatalf("Upsert target %s: %v", d.Title, err)
}
}
// Wikilinks src → each target.
for _, toID := range []string{shared.ID, personalSelf.ID, personalOther.ID, teamOther.ID} {
if err := vs.CreateLink(ctxSelf, &store.VaultLink{
FromDocID: src.ID,
ToDocID: toID,
LinkType: "wikilink",
}); err != nil {
t.Fatalf("CreateLink → %s: %v", toID, err)
}
}
tool := tools.NewVaultReadTool()
tool.SetVaultStore(vs)
tool.SetWorkspace(ws)
res := tool.Execute(ctxSelf, map[string]any{"doc_id": src.ID})
if res.IsError {
t.Fatalf("vault_read error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "## Links") {
t.Fatalf("missing Links section: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "SharedTgt") {
t.Fatalf("expected shared target in footer: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "SelfTgt") {
t.Fatalf("expected personal-self target in footer: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "PersonalOtherTgt") {
t.Fatalf("personal-other leaked: %s", res.ForLLM)
}
if strings.Contains(res.ForLLM, "TeamOtherTgt") {
t.Fatalf("team-other leaked: %s", res.ForLLM)
}
}