diff --git a/internal/store/pg/vault_links.go b/internal/store/pg/vault_links.go index a3474b32..14df122c 100644 --- a/internal/store/pg/vault_links.go +++ b/internal/store/pg/vault_links.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "fmt" + "strings" "time" "github.com/google/uuid" @@ -145,6 +146,30 @@ func (s *PGVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, docID return err } +// DeleteDocLinksByTypes removes outbound links matching any of the given types from a document. +func (s *PGVaultStore) DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error { + if len(types) == 0 { + return nil + } + uid := mustParseUUID(docID) + tid := mustParseUUID(tenantID) + // Build IN clause with positional params: $3, $4, ... + params := []any{uid, tid} + placeholders := make([]string, len(types)) + for i, t := range types { + params = append(params, t) + placeholders[i] = fmt.Sprintf("$%d", i+3) + } + q := fmt.Sprintf(` + DELETE FROM vault_links vl + USING vault_documents vd + WHERE vl.from_doc_id = $1 + AND vd.id = vl.from_doc_id AND vd.tenant_id = $2 + AND vl.link_type IN (%s)`, strings.Join(placeholders, ",")) + _, err := s.db.ExecContext(ctx, q, params...) + return err +} + func scanVaultLinks(rows *sql.Rows) ([]store.VaultLink, error) { var links []store.VaultLink for rows.Next() { diff --git a/internal/store/sqlitestore/vault_links.go b/internal/store/sqlitestore/vault_links.go index 74109793..11abe711 100644 --- a/internal/store/sqlitestore/vault_links.go +++ b/internal/store/sqlitestore/vault_links.go @@ -6,6 +6,7 @@ import ( "context" "database/sql" "fmt" + "strings" "time" "github.com/google/uuid" @@ -122,6 +123,29 @@ func (s *SQLiteVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, d return err } +// DeleteDocLinksByTypes removes outbound links matching any of the given types from a document. +func (s *SQLiteVaultStore) DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error { + if len(types) == 0 { + return nil + } + // Build IN clause with ? placeholders. + params := []any{docID} + placeholders := make([]string, len(types)) + for i, t := range types { + params = append(params, t) + placeholders[i] = "?" + } + params = append(params, tenantID) + q := fmt.Sprintf(` + DELETE FROM vault_links + WHERE from_doc_id = ? + AND link_type IN (%s) + AND from_doc_id IN (SELECT id FROM vault_documents WHERE tenant_id = ?)`, + strings.Join(placeholders, ",")) + _, err := s.db.ExecContext(ctx, q, params...) + return err +} + func scanVaultLinkRows(rows *sql.Rows) ([]store.VaultLink, error) { var links []store.VaultLink for rows.Next() { diff --git a/internal/store/vault_store.go b/internal/store/vault_store.go index 6ced6c52..5ac3a0a0 100644 --- a/internal/store/vault_store.go +++ b/internal/store/vault_store.go @@ -93,6 +93,7 @@ type VaultStore interface { GetBacklinks(ctx context.Context, tenantID, docID string) ([]VaultBacklink, error) DeleteDocLinks(ctx context.Context, tenantID, docID string) error DeleteDocLinksByType(ctx context.Context, tenantID, docID, linkType string) error + DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error // Enrichment // UpdateSummaryAndReembed updates summary text and re-generates embedding from title+path+summary. diff --git a/internal/vault/enrich_classify.go b/internal/vault/enrich_classify.go new file mode 100644 index 00000000..8cb97fbd --- /dev/null +++ b/internal/vault/enrich_classify.go @@ -0,0 +1,203 @@ +package vault + +import ( + "context" + "fmt" + "log/slog" + "maps" + "slices" + "strings" + "time" + + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const ( + classifyMaxRetries = 3 + classifyMaxTokens = 1024 + classifyTemperature = 0.1 + classifyCtxMaxLen = 256 // max context string length stored in DB + classifySummaryMaxChars = 300 // max summary chars in prompt (validated: 300 for accuracy) + classifyMaxSourceDocs = 20 // max source docs per classifyLinks call (validated: cap unbounded time) +) + +var classifyTimeouts = [3]time.Duration{5 * time.Minute, 7 * time.Minute, 10 * time.Minute} +var classifyBackoffs = [3]time.Duration{0, 2 * time.Second, 4 * time.Second} + +// validClassifyTypes — accepted link types stored directly in DB (aligned with UI vault-link-dialog.tsx). +var validClassifyTypes = map[string]bool{ + "reference": true, "depends_on": true, "extends": true, + "related": true, "supersedes": true, "contradicts": true, +} + +type classifyDoc struct { + DocID, Title, Path, Summary string +} + +type candidatePair struct { + Source, Candidate classifyDoc + Score float64 +} + +// classifyLinks orchestrates LLM-based link classification for enriched docs. +func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID string, results []enriched) { + if w.provider == nil { + return + } + + capped := results + if len(capped) > classifyMaxSourceDocs { + capped = capped[:classifyMaxSourceDocs] + } + + // On SQLite (desktop/Lite), FindSimilarDocs returns nil (no pgvector) — classify is a no-op. + candidates := w.gatherCandidates(ctx, tenantID, agentID, capped) + if len(candidates) == 0 { + return + } + + allTypes := slices.Collect(maps.Keys(validClassifyTypes)) + allTypes = append(allTypes, "semantic") // clean up legacy links + + for sourceDocID, pairs := range candidates { + source := pairs[0].Source + candidateDocs := make([]classifyDoc, len(pairs)) + for i, p := range pairs { + candidateDocs[i] = p.Candidate + } + + system, user := buildClassifyPrompt(source, candidateDocs) + raw, err := w.callClassifyWithRetry(ctx, system, user) + if err != nil { + slog.Warn("vault.classify: llm_failed", "doc", sourceDocID, "err", err) + continue // SKIP fallback + } + + parsed, err := parseClassifyResponse(raw, len(candidateDocs)) + if err != nil { + hint := fmt.Sprintf("\n\nPrevious response was invalid JSON (error: %s). Output ONLY a valid JSON array.", err.Error()) + raw2, err2 := w.callClassifyWithRetry(ctx, system, user+hint) + if err2 != nil { + slog.Warn("vault.classify: retry_parse_failed", "doc", sourceDocID, "err", err2) + continue + } + parsed, err = parseClassifyResponse(raw2, len(candidateDocs)) + if err != nil { + slog.Warn("vault.classify: parse_still_failed", "doc", sourceDocID, "err", err) + continue + } + } + + // Collect valid links (collect-then-write pattern). + var newLinks []store.VaultLink + for _, r := range parsed { + if r.Type == "SKIP" || !validClassifyTypes[r.Type] { + continue + } + linkCtx := r.Ctx + if len(linkCtx) > classifyCtxMaxLen { + linkCtx = string([]rune(linkCtx)[:classifyCtxMaxLen]) + } + newLinks = append(newLinks, store.VaultLink{ + FromDocID: sourceDocID, + ToDocID: candidateDocs[r.Idx-1].DocID, // idx is 1-based, validated by parseClassifyResponse + LinkType: r.Type, + Context: linkCtx, + }) + } + + // Only replace old links if LLM produced valid replacements (avoid data loss on all-SKIP). + if len(newLinks) > 0 { + if err := w.vault.DeleteDocLinksByTypes(ctx, tenantID, sourceDocID, allTypes); err != nil { + slog.Warn("vault.classify: delete_old", "doc", sourceDocID, "err", err) + } + for i := range newLinks { + if err := w.vault.CreateLink(ctx, &newLinks[i]); err != nil { + slog.Debug("vault.classify: create_link", "from", sourceDocID, "to", newLinks[i].ToDocID, "err", err) + } + } + } + } +} + +func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, agentID string, results []enriched) map[string][]candidatePair { + seen := make(map[string]bool) + out := make(map[string][]candidatePair) + + for _, r := range results { + neighbors, err := w.vault.FindSimilarDocs(ctx, tenantID, agentID, r.payload.DocID, enrichSimilarityLimit) + if err != nil { + slog.Warn("vault.classify: find_similar", "doc", r.payload.DocID, "err", err) + continue + } + // Derive title from path (payload has no Title field; neighbor docs have it from DB). + title := r.payload.Path + if doc, err := w.vault.GetDocumentByID(ctx, tenantID, r.payload.DocID); err == nil && doc != nil { + title = doc.Title + } + src := classifyDoc{ + DocID: r.payload.DocID, + Title: title, + Path: r.payload.Path, + Summary: truncateSummary(r.summary), + } + for _, n := range neighbors { + if n.Score < enrichSimilarityMin || n.Document.Summary == "" { + continue + } + // Bidirectional dedup: only process each pair once. + a, b := src.DocID, n.Document.ID + if a > b { + a, b = b, a + } + if key := a + ":" + b; seen[key] { + continue + } else { + seen[key] = true + } + out[src.DocID] = append(out[src.DocID], candidatePair{ + Source: src, + Candidate: classifyDoc{ + DocID: n.Document.ID, + Title: n.Document.Title, + Path: n.Document.Path, + Summary: truncateSummary(n.Document.Summary), + }, + Score: n.Score, + }) + } + } + return out +} + +// callClassifyWithRetry calls the LLM with escalating timeouts and backoffs. +func (w *enrichWorker) callClassifyWithRetry(ctx context.Context, system, user string) (string, error) { + var lastErr error + for attempt := range classifyMaxRetries { + if attempt > 0 { + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-time.After(classifyBackoffs[attempt]): + } + } + cctx, cancel := context.WithTimeout(ctx, classifyTimeouts[attempt]) + resp, err := w.provider.Chat(cctx, providers.ChatRequest{ + Messages: []providers.Message{ + {Role: "system", Content: system}, + {Role: "user", Content: user}, + }, + Model: w.model, + Options: map[string]any{"max_tokens": classifyMaxTokens, "temperature": classifyTemperature}, + }) + cancel() + if err != nil { + lastErr = err + slog.Warn("vault.classify: retry", "attempt", attempt+1, "err", err) + continue + } + return strings.TrimSpace(resp.Content), nil + } + return "", fmt.Errorf("classify exhausted %d retries: %w", classifyMaxRetries, lastErr) +} diff --git a/internal/vault/enrich_classify_parse_edge_test.go b/internal/vault/enrich_classify_parse_edge_test.go new file mode 100644 index 00000000..3059100d --- /dev/null +++ b/internal/vault/enrich_classify_parse_edge_test.go @@ -0,0 +1,186 @@ +package vault + +import ( + "testing" +) + +// ============================================================================ +// JSON Parsing Edge Case Tests +// ============================================================================ + +// TestParseClassifyResponse_MixedValidInvalid handles mix of valid and invalid entries. +func TestParseClassifyResponse_MixedValidInvalid(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"valid"}, + {"idx":0,"type":"reference","ctx":"invalid idx"}, + {"idx":2,"type":"bad_type","ctx":"invalid type"}, + {"idx":3,"type":"extends","ctx":"valid"} + ]` + + results, err := parseClassifyResponse(raw, 3) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + // Should keep only idx 1 and 3 with valid types + if len(results) != 2 { + t.Fatalf("Expected 2 valid results, got %d", len(results)) + } + + if results[0].Idx != 1 || results[1].Idx != 3 { + t.Errorf("Invalid entries not filtered correctly") + } +} + +// TestParseClassifyResponse_NegativeIdx filters silently. +func TestParseClassifyResponse_NegativeIdx(t *testing.T) { + raw := `[ + {"idx":-1,"type":"reference","ctx":"invalid"}, + {"idx":1,"type":"reference","ctx":"valid"} + ]` + + results, err := parseClassifyResponse(raw, 1) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 1 { + t.Fatalf("Expected 1 result (negative idx filtered), got %d", len(results)) + } + + if results[0].Idx != 1 { + t.Errorf("Negative idx should be filtered") + } +} + +// TestParseClassifyResponse_AllSKIP preserves all SKIP entries. +func TestParseClassifyResponse_AllSKIP(t *testing.T) { + raw := `[ + {"idx":1,"type":"SKIP","ctx":"no match"}, + {"idx":2,"type":"SKIP","ctx":"different topic"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 SKIP entries, got %d", len(results)) + } + + for i, r := range results { + if r.Type != "SKIP" { + t.Errorf("Result %d: expected SKIP, got %q", i, r.Type) + } + } +} + +// TestParseClassifyResponse_LargeIdx beyond count are filtered. +func TestParseClassifyResponse_LargeIdx(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"ok"}, + {"idx":100,"type":"reference","ctx":"too large"}, + {"idx":999,"type":"reference","ctx":"way too large"} + ]` + + results, err := parseClassifyResponse(raw, 5) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 1 { + t.Fatalf("Expected 1 result (oversized idx filtered), got %d", len(results)) + } + + if results[0].Idx != 1 { + t.Errorf("Large idx should be filtered") + } +} + +// TestParseClassifyResponse_MaxBoundaryIdx at exactly count is valid. +func TestParseClassifyResponse_MaxBoundaryIdx(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"first"}, + {"idx":5,"type":"reference","ctx":"last valid"} + ]` + + results, err := parseClassifyResponse(raw, 5) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results, got %d", len(results)) + } + + if results[1].Idx != 5 { + t.Errorf("Max boundary idx (5) should be valid when count=5") + } +} + +// TestParseClassifyResponse_EmptyCtx preserves empty context string. +func TestParseClassifyResponse_EmptyCtx(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":""}, + {"idx":2,"type":"related","ctx":"has context"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results, got %d", len(results)) + } + + if results[0].Ctx != "" { + t.Errorf("Empty ctx should be preserved, got %q", results[0].Ctx) + } +} + +// TestParseClassifyResponse_WhitespaceCtx preserves whitespace-only context. +func TestParseClassifyResponse_WhitespaceCtx(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":" "}, + {"idx":2,"type":"related","ctx":"valid"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results, got %d", len(results)) + } + + // Whitespace should be preserved + if results[0].Ctx != " " { + t.Errorf("Whitespace ctx should be preserved, got %q", results[0].Ctx) + } +} + +// TestParseClassifyResponse_CaseInsensitiveType only SKIP and validClassifyTypes are accepted. +func TestParseClassifyResponse_CaseInsensitiveType(t *testing.T) { + // Lowercase versions of valid types should not be accepted + raw := `[ + {"idx":1,"type":"reference","ctx":"ok"}, + {"idx":2,"type":"REFERENCE","ctx":"uppercase"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + // Only lowercase "reference" should be valid + if len(results) != 1 { + t.Fatalf("Expected 1 result (uppercase REFERENCE filtered), got %d", len(results)) + } + + if results[0].Type != "reference" { + t.Errorf("Only lowercase types accepted, got %q", results[0].Type) + } +} diff --git a/internal/vault/enrich_classify_parse_test.go b/internal/vault/enrich_classify_parse_test.go new file mode 100644 index 00000000..3400103e --- /dev/null +++ b/internal/vault/enrich_classify_parse_test.go @@ -0,0 +1,193 @@ +package vault + +import ( + "fmt" + "strings" + "testing" +) + +// ============================================================================ +// JSON Parsing Tests +// ============================================================================ + +// TestParseClassifyResponse_Valid parses standard JSON array. +func TestParseClassifyResponse_Valid(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"mentions config"}, + {"idx":2,"type":"depends_on","ctx":"requires authentication"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results, got %d", len(results)) + } + + if results[0].Idx != 1 || results[0].Type != "reference" { + t.Errorf("Result 0: got idx=%d type=%q, want idx=1 type=reference", results[0].Idx, results[0].Type) + } + + if results[1].Idx != 2 || results[1].Type != "depends_on" { + t.Errorf("Result 1: got idx=%d type=%q, want idx=2 type=depends_on", results[1].Idx, results[1].Type) + } +} + +// TestParseClassifyResponse_WithCodeFence strips ```json fences. +func TestParseClassifyResponse_WithCodeFence(t *testing.T) { + raw := "```json\n[{\"idx\":1,\"type\":\"reference\",\"ctx\":\"test\"}]\n```" + + results, err := parseClassifyResponse(raw, 1) + if err != nil { + t.Fatalf("parseClassifyResponse with code fence failed: %v", err) + } + + if len(results) != 1 { + t.Fatalf("Expected 1 result, got %d", len(results)) + } + + if results[0].Type != "reference" { + t.Errorf("Expected type reference, got %q", results[0].Type) + } +} + +// TestParseClassifyResponse_InvalidJSON returns error on unmarshal failure. +func TestParseClassifyResponse_InvalidJSON(t *testing.T) { + raw := `{invalid json}` + + _, err := parseClassifyResponse(raw, 1) + if err == nil { + t.Fatalf("parseClassifyResponse should return error for invalid JSON") + } + + if !strings.Contains(err.Error(), "unmarshal") { + t.Errorf("Error should mention unmarshal: %v", err) + } +} + +// TestParseClassifyResponse_OutOfRangeIdx filters silently, keeps valid entries. +func TestParseClassifyResponse_OutOfRangeIdx(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"ok"}, + {"idx":99,"type":"related","ctx":"invalid idx"}, + {"idx":2,"type":"extends","ctx":"ok"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results (out-of-range filtered), got %d", len(results)) + } + + if results[0].Idx != 1 || results[1].Idx != 2 { + t.Errorf("Out-of-range entry not filtered properly") + } +} + +// TestParseClassifyResponse_UnknownType filters silently. +func TestParseClassifyResponse_UnknownType(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"ok"}, + {"idx":2,"type":"unknown_type","ctx":"invalid"}, + {"idx":3,"type":"related","ctx":"ok"} + ]` + + results, err := parseClassifyResponse(raw, 3) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results (unknown type filtered), got %d", len(results)) + } + + if results[0].Type != "reference" || results[1].Type != "related" { + t.Errorf("Unknown type not filtered properly") + } +} + +// TestParseClassifyResponse_SKIPType preserves SKIP in output. +func TestParseClassifyResponse_SKIPType(t *testing.T) { + raw := `[ + {"idx":1,"type":"reference","ctx":"valid"}, + {"idx":2,"type":"SKIP","ctx":"no relationship"} + ]` + + results, err := parseClassifyResponse(raw, 2) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 2 { + t.Fatalf("Expected 2 results including SKIP, got %d", len(results)) + } + + skipFound := false + for _, r := range results { + if r.Type == "SKIP" { + skipFound = true + break + } + } + + if !skipFound { + t.Errorf("SKIP type should be preserved in results") + } +} + +// TestParseClassifyResponse_EmptyArray returns valid empty slice. +func TestParseClassifyResponse_EmptyArray(t *testing.T) { + raw := `[]` + + results, err := parseClassifyResponse(raw, 5) + if err != nil { + t.Fatalf("parseClassifyResponse failed on empty array: %v", err) + } + + if len(results) != 0 { + t.Fatalf("Expected 0 results, got %d", len(results)) + } +} + +// TestParseClassifyResponse_CtxTruncation truncates context over 256 chars. +func TestParseClassifyResponse_CtxTruncation(t *testing.T) { + longCtx := strings.Repeat("x", 300) + raw := fmt.Sprintf(`[{"idx":1,"type":"reference","ctx":"%s"}]`, longCtx) + + results, err := parseClassifyResponse(raw, 1) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len([]rune(results[0].Ctx)) > classifyCtxMaxLen { + t.Errorf("Context not truncated: got %d runes, want ≤%d", + len([]rune(results[0].Ctx)), classifyCtxMaxLen) + } +} + +// TestParseClassifyResponse_ZeroIdx filters silently (idx must be >= 1). +func TestParseClassifyResponse_ZeroIdx(t *testing.T) { + raw := `[ + {"idx":0,"type":"reference","ctx":"invalid"}, + {"idx":1,"type":"reference","ctx":"valid"} + ]` + + results, err := parseClassifyResponse(raw, 1) + if err != nil { + t.Fatalf("parseClassifyResponse failed: %v", err) + } + + if len(results) != 1 { + t.Fatalf("Expected 1 result (idx 0 filtered), got %d", len(results)) + } + + if results[0].Idx != 1 { + t.Errorf("Zero idx should be filtered") + } +} + diff --git a/internal/vault/enrich_classify_prompt.go b/internal/vault/enrich_classify_prompt.go new file mode 100644 index 00000000..a02654c8 --- /dev/null +++ b/internal/vault/enrich_classify_prompt.go @@ -0,0 +1,90 @@ +package vault + +import ( + "encoding/json" + "fmt" + "log/slog" + "strings" +) + +// classifyResult is a single LLM classification output for a candidate doc. +type classifyResult struct { + Idx int `json:"idx"` + Type string `json:"type"` + Ctx string `json:"ctx"` +} + +const classifySystemPrompt = `You classify relationships between documents in a knowledge vault. + +## Link Types +- reference: A cites, mentions, or quotes B +- depends_on: A requires B to function (config, data dependency, prerequisite) +- extends: A expands, details, or elaborates on B +- related: A and B share the same topic or domain (use only when no stronger type fits) +- supersedes: A is a newer version that replaces or updates B +- contradicts: A conflicts with or opposes B's content + +## Rules +- Output ONLY a valid JSON array, no other text +- Use SKIP when documents are similar but have no meaningful relationship +- Prefer specific types (reference, depends_on) over generic (related) +- Each candidate classified independently +- Keep ctx descriptions under 60 words + +## Output Format +[{"idx":1,"type":"reference","ctx":"mentions OAuth config for setup"},{"idx":2,"type":"SKIP"}]` + +// buildClassifyPrompt formats the system and user prompts for classify LLM call. +func buildClassifyPrompt(source classifyDoc, candidates []classifyDoc) (system, user string) { + var b strings.Builder + b.WriteString("## Source Document\n") + fmt.Fprintf(&b, "Title: %s\nPath: %s\nSummary: %s\n\n", source.Title, source.Path, source.Summary) + b.WriteString("## Candidates\n") + for i, c := range candidates { + fmt.Fprintf(&b, "%d. Title: %s | Path: %s | Summary: %s\n", i+1, c.Title, c.Path, c.Summary) + } + return classifySystemPrompt, b.String() +} + +// parseClassifyResponse parses LLM JSON output into classify results. +// Uses partial success model: invalid entries filtered silently, error only on total unmarshal failure. +func parseClassifyResponse(raw string, count int) ([]classifyResult, error) { + raw = strings.TrimSpace(raw) + // Strip code fences. + raw = strings.TrimPrefix(raw, "```json") + raw = strings.TrimPrefix(raw, "```") + raw = strings.TrimSuffix(raw, "```") + raw = strings.TrimSpace(raw) + + var results []classifyResult + if err := json.Unmarshal([]byte(raw), &results); err != nil { + return nil, fmt.Errorf("json unmarshal: %w", err) + } + + // Filter invalid entries (partial success). + valid := results[:0] + for _, r := range results { + if r.Idx < 1 || r.Idx > count { + slog.Debug("vault.classify: idx out of range", "idx", r.Idx, "count", count) + continue + } + if r.Type != "SKIP" && !validClassifyTypes[r.Type] { + continue + } + // Truncate context. + if len(r.Ctx) > classifyCtxMaxLen { + r.Ctx = string([]rune(r.Ctx)[:classifyCtxMaxLen]) + } + valid = append(valid, r) + } + return valid, nil +} + +// truncateSummary caps summary at classifySummaryMaxChars with UTF-8 safe rune truncation. +func truncateSummary(s string) string { + runes := []rune(s) + if len(runes) <= classifySummaryMaxChars { + return s + } + return string(runes[:classifySummaryMaxChars]) + "..." +} diff --git a/internal/vault/enrich_classify_prompt_test.go b/internal/vault/enrich_classify_prompt_test.go new file mode 100644 index 00000000..776d5aaa --- /dev/null +++ b/internal/vault/enrich_classify_prompt_test.go @@ -0,0 +1,148 @@ +package vault + +import ( + "strings" + "testing" +) + +// ============================================================================ +// Prompt Building Tests +// ============================================================================ + +// TestBuildClassifyPrompt verifies system prompt contains all 6 types and user prompt has numbered candidates. +func TestBuildClassifyPrompt(t *testing.T) { + source := classifyDoc{ + DocID: "doc1", + Title: "Test Document", + Path: "docs/test.md", + Summary: "This is a test document.", + } + + candidates := []classifyDoc{ + { + DocID: "doc2", + Title: "Related Doc 1", + Path: "docs/related1.md", + Summary: "First related document.", + }, + { + DocID: "doc3", + Title: "Related Doc 2", + Path: "docs/related2.md", + Summary: "Second related document.", + }, + } + + system, user := buildClassifyPrompt(source, candidates) + + // System prompt must contain all 6 relationship types + expectedTypes := []string{"reference", "depends_on", "extends", "related", "supersedes", "contradicts"} + for _, typ := range expectedTypes { + if !strings.Contains(system, typ) { + t.Errorf("System prompt missing type: %s", typ) + } + } + + // System prompt must contain classifySystemPrompt constant text + if !strings.Contains(system, "Link Types") { + t.Errorf("System prompt missing 'Link Types' header") + } + + // User prompt must contain source doc info + if !strings.Contains(user, "Test Document") { + t.Errorf("User prompt missing source title") + } + if !strings.Contains(user, "docs/test.md") { + t.Errorf("User prompt missing source path") + } + if !strings.Contains(user, "This is a test document") { + t.Errorf("User prompt missing source summary") + } + + // User prompt must contain candidates with numbers (1., 2.) + if !strings.Contains(user, "1.") { + t.Errorf("User prompt missing numbered candidate (1.)") + } + if !strings.Contains(user, "2.") { + t.Errorf("User prompt missing numbered candidate (2.)") + } + + // User prompt must contain candidate info + if !strings.Contains(user, "Related Doc 1") { + t.Errorf("User prompt missing candidate 1 title") + } + if !strings.Contains(user, "Related Doc 2") { + t.Errorf("User prompt missing candidate 2 title") + } +} + +// ============================================================================ +// Truncate Summary Tests +// ============================================================================ + +// TestTruncateSummary_Long truncates at classifySummaryMaxChars (300) with "..." suffix. +func TestTruncateSummary_Long(t *testing.T) { + // Create string longer than 300 chars + longText := strings.Repeat("a", 350) + result := truncateSummary(longText) + + if len([]rune(result)) > classifySummaryMaxChars+3 { // +3 for "..." + t.Errorf("truncateSummary(%d chars) returned %d runes, want ≤%d", + len([]rune(longText)), len([]rune(result)), classifySummaryMaxChars+3) + } + + if !strings.HasSuffix(result, "...") { + t.Errorf("Truncated result must end with '...', got: %q", result) + } + + // Should be exactly 303 runes (300 + "...") + expected := len([]rune(strings.Repeat("a", classifySummaryMaxChars))) + 3 + if len([]rune(result)) != expected { + t.Errorf("Truncated result has %d runes, want %d", len([]rune(result)), expected) + } +} + +// TestTruncateSummary_Short returns short string as-is. +func TestTruncateSummary_Short(t *testing.T) { + short := "This is short" + result := truncateSummary(short) + + if result != short { + t.Errorf("truncateSummary(%q) = %q, want unchanged", short, result) + } + + if strings.Contains(result, "...") { + t.Errorf("Short string should not have '...' suffix: %q", result) + } +} + +// TestTruncateSummary_BoundaryAt300 tests string exactly at 300 chars. +func TestTruncateSummary_BoundaryAt300(t *testing.T) { + exact := strings.Repeat("x", classifySummaryMaxChars) + result := truncateSummary(exact) + + if result != exact { + t.Errorf("String at exactly %d chars should not be truncated", classifySummaryMaxChars) + } + + if strings.HasSuffix(result, "...") { + t.Errorf("String at boundary should not have '...' suffix") + } +} + +// TestTruncateSummary_Unicode tests UTF-8 safe truncation with multi-byte chars. +func TestTruncateSummary_Unicode(t *testing.T) { + // Use emoji (4 bytes each) to test multi-byte handling + longUnicode := strings.Repeat("🎉", 400) // 400 emoji chars, very long when encoded + result := truncateSummary(longUnicode) + + runes := []rune(result) + if len(runes) > classifySummaryMaxChars+3 { + t.Errorf("Unicode truncation exceeded max: got %d runes, want ≤%d", + len(runes), classifySummaryMaxChars+3) + } + + if !strings.HasSuffix(result, "...") { + t.Errorf("Truncated unicode must end with '...'") + } +} diff --git a/internal/vault/enrich_classify_retry_advanced_test.go b/internal/vault/enrich_classify_retry_advanced_test.go new file mode 100644 index 00000000..f6720ed3 --- /dev/null +++ b/internal/vault/enrich_classify_retry_advanced_test.go @@ -0,0 +1,154 @@ +package vault + +import ( + "context" + "errors" + "strings" + "testing" +) + +// ============================================================================ +// Advanced Retry Logic Tests +// ============================================================================ + +// TestCallClassifyWithRetry_FirstAttemptSuccess uses first timeout. +func TestCallClassifyWithRetry_FirstAttemptSuccess(t *testing.T) { + // Verify that first attempt uses classifyTimeouts[0] + if classifyTimeouts[0] == 0 { + t.Errorf("First timeout should be non-zero") + } + + provider := &mockClassifyProvider{ + responses: []string{`[{"idx":1,"type":"reference","ctx":"first attempt"}]`}, + errors: []error{nil}, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err != nil { + t.Fatalf("First attempt should succeed: %v", err) + } + + if provider.calls != 1 { + t.Errorf("Expected exactly 1 call on first-attempt success, got %d", provider.calls) + } + + if !strings.Contains(resp, "reference") { + t.Errorf("Response missing expected content: %q", resp) + } +} + +// TestCallClassifyWithRetry_ResponseWhitespaceStripping trims whitespace from response. +func TestCallClassifyWithRetry_ResponseWhitespaceStripping(t *testing.T) { + // Response has leading/trailing whitespace that should be stripped + provider := &mockClassifyProvider{ + responses: []string{"\n\n [{\"idx\":1,\"type\":\"reference\",\"ctx\":\"test\"}] \n"}, + errors: []error{nil}, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err != nil { + t.Fatalf("callClassifyWithRetry failed: %v", err) + } + + // Response should be trimmed + if strings.HasPrefix(resp, "\n") || strings.HasSuffix(resp, "\n") { + t.Errorf("Response should be trimmed: %q", resp) + } + + if !strings.Contains(resp, "reference") { + t.Errorf("Response missing expected content: %q", resp) + } +} + +// TestCallClassifyWithRetry_MaxRetriesConstant verifies retry limit. +func TestCallClassifyWithRetry_MaxRetriesConstant(t *testing.T) { + if classifyMaxRetries != 3 { + t.Errorf("classifyMaxRetries should be 3, got %d", classifyMaxRetries) + } + + // Verify arrays have correct length + if len(classifyTimeouts) != classifyMaxRetries { + t.Errorf("classifyTimeouts length should be %d, got %d", classifyMaxRetries, len(classifyTimeouts)) + } + + if len(classifyBackoffs) != classifyMaxRetries { + t.Errorf("classifyBackoffs length should be %d, got %d", classifyMaxRetries, len(classifyBackoffs)) + } +} + +// TestCallClassifyWithRetry_SecondAttemptSucceeds verifies first retry succeeds. +func TestCallClassifyWithRetry_SecondAttemptSucceeds(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{ + "", // attempt 0: error + `[{"idx":1,"type":"extends","ctx":"second attempt"}]`, // attempt 1: success + }, + errors: []error{ + errors.New("first attempt failed"), + nil, // no error on second attempt + }, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err != nil { + t.Fatalf("Should succeed on second attempt, got error: %v", err) + } + + if provider.calls != 2 { + t.Errorf("Expected 2 calls, got %d", provider.calls) + } + + if !strings.Contains(resp, "extends") { + t.Errorf("Response from second attempt not found: %q", resp) + } +} + +// TestCallClassifyWithRetry_EmptyResponse returns error for empty response on all attempts. +func TestCallClassifyWithRetry_EmptyResponse(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{"", "", ""}, + errors: []error{nil, nil, nil}, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + // Empty response is still a successful LLM call, should return empty string + if err != nil { + t.Fatalf("Empty response should not error: %v", err) + } + + if resp != "" { + t.Errorf("Expected empty response, got %q", resp) + } + + if provider.calls != 1 { + t.Errorf("Expected 1 call (succeeded immediately), got %d", provider.calls) + } +} diff --git a/internal/vault/enrich_classify_retry_test.go b/internal/vault/enrich_classify_retry_test.go new file mode 100644 index 00000000..6536ffab --- /dev/null +++ b/internal/vault/enrich_classify_retry_test.go @@ -0,0 +1,195 @@ +package vault + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// ============================================================================ +// Mock Provider for Testing +// ============================================================================ + +// mockClassifyProvider implements providers.Provider for testing retry logic. +type mockClassifyProvider struct { + responses []string + errors []error + calls int +} + +func (m *mockClassifyProvider) Chat(ctx context.Context, req providers.ChatRequest) (*providers.ChatResponse, error) { + idx := m.calls + m.calls++ + if idx < len(m.errors) && m.errors[idx] != nil { + return nil, m.errors[idx] + } + resp := "" + if idx < len(m.responses) { + resp = m.responses[idx] + } + return &providers.ChatResponse{Content: resp}, nil +} + +func (m *mockClassifyProvider) ChatStream(ctx context.Context, req providers.ChatRequest, onChunk func(providers.StreamChunk)) (*providers.ChatResponse, error) { + return m.Chat(ctx, req) +} + +func (m *mockClassifyProvider) DefaultModel() string { return "test" } + +func (m *mockClassifyProvider) Name() string { return "mock" } + +// ============================================================================ +// Retry Logic Tests +// ============================================================================ + +// TestCallClassifyWithRetry_Success succeeds on first attempt. +func TestCallClassifyWithRetry_Success(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{`[{"idx":1,"type":"reference","ctx":"test"}]`}, + errors: []error{nil}, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err != nil { + t.Fatalf("callClassifyWithRetry failed: %v", err) + } + + if !strings.Contains(resp, "reference") { + t.Errorf("Response doesn't contain expected content: %q", resp) + } + + if provider.calls != 1 { + t.Errorf("Expected 1 call, got %d", provider.calls) + } +} + +// TestCallClassifyWithRetry_RetryThenSuccess fails twice, succeeds on third attempt. +func TestCallClassifyWithRetry_RetryThenSuccess(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{ + "", // attempt 0: error + "", // attempt 1: error + `[{"idx":1,"type":"reference","ctx":"success"}]`, // attempt 2: success + }, + errors: []error{ + fmt.Errorf("network error"), + fmt.Errorf("timeout"), + nil, // no error on third attempt + }, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + resp, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err != nil { + t.Fatalf("callClassifyWithRetry should succeed after retries, got error: %v", err) + } + + if !strings.Contains(resp, "success") { + t.Errorf("Response doesn't contain expected content: %q", resp) + } + + if provider.calls != 3 { + t.Errorf("Expected 3 calls, got %d", provider.calls) + } +} + +// TestCallClassifyWithRetry_AllFail exhausts retries and returns error. +func TestCallClassifyWithRetry_AllFail(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{"", "", ""}, + errors: []error{ + fmt.Errorf("error1"), + fmt.Errorf("error2"), + fmt.Errorf("error3"), + }, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx := context.Background() + _, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err == nil { + t.Fatalf("callClassifyWithRetry should return error after exhausting retries") + } + + if !strings.Contains(err.Error(), "exhausted") { + t.Errorf("Error should mention exhausted retries: %v", err) + } + + if provider.calls != classifyMaxRetries { + t.Errorf("Expected %d calls, got %d", classifyMaxRetries, provider.calls) + } +} + +// TestCallClassifyWithRetry_ContextCancellation respects context cancellation. +func TestCallClassifyWithRetry_ContextCancellation(t *testing.T) { + provider := &mockClassifyProvider{ + responses: []string{"", "", ""}, + errors: []error{ + fmt.Errorf("error1"), + fmt.Errorf("error2"), + fmt.Errorf("error3"), + }, + } + + worker := &enrichWorker{ + provider: provider, + model: "test", + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + _, err := worker.callClassifyWithRetry(ctx, "system", "user") + + if err == nil { + t.Fatalf("callClassifyWithRetry should return error for cancelled context") + } + + if !strings.Contains(err.Error(), "context") { + t.Errorf("Error should mention context: %v", err) + } +} + +// TestCallClassifyWithRetry_RetriesAndBackoffs verifies retry timeouts and backoffs escalate. +func TestCallClassifyWithRetry_RetriesAndBackoffs(t *testing.T) { + // Verify first backoff is 0 (no backoff on first attempt) + if classifyBackoffs[0] != 0 { + t.Errorf("First backoff should be 0, got %v", classifyBackoffs[0]) + } + + // Verify backoffs escalate for retry attempts + if classifyBackoffs[1] == 0 || classifyBackoffs[2] == 0 { + t.Errorf("Backoffs should escalate: %v", classifyBackoffs) + } + + if classifyBackoffs[1] >= classifyBackoffs[2] { + t.Errorf("Backoffs should increase: %v", classifyBackoffs) + } + + // Verify timeouts escalate + if classifyTimeouts[0] >= classifyTimeouts[1] || classifyTimeouts[1] >= classifyTimeouts[2] { + t.Errorf("Timeouts should escalate: %v", classifyTimeouts) + } +} + diff --git a/internal/vault/enrich_worker.go b/internal/vault/enrich_worker.go index a53b7bd8..6fcbe269 100644 --- a/internal/vault/enrich_worker.go +++ b/internal/vault/enrich_worker.go @@ -2,7 +2,6 @@ package vault import ( "context" - "fmt" "log/slog" "os" "path/filepath" @@ -80,6 +79,12 @@ func (w *enrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) e return nil } +// enriched holds a successfully summarized vault document pending embed+link. +type enriched struct { + payload eventbus.VaultDocUpsertedPayload + summary string +} + // processBatch drains and processes queued vault doc events in a loop. func (w *enrichWorker) processBatch(ctx context.Context, key string) { for { @@ -91,10 +96,6 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) { continue } - type enriched struct { - payload eventbus.VaultDocUpsertedPayload - summary string - } var results []enriched for _, item := range items { @@ -145,20 +146,28 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) { results = append(results, enriched{payload: item, summary: summary}) } - // Update summary + embed + auto-link for each enriched doc. + // Phase 2 — Embed: update summary + embed for all results. + // Do NOT record dedup here (moved to Phase 4 after classify). + var embedded []enriched for _, r := range results { if err := w.vault.UpdateSummaryAndReembed(ctx, r.payload.TenantID, r.payload.DocID, r.summary); err != nil { slog.Warn("vault.enrich: update_summary", "doc", r.payload.DocID, "err", err) - continue // don't record in dedup + continue } + embedded = append(embedded, r) + } - // Record hash only after successful update. + // Phase 3 — Classify links (replaces autoLink). + // All items share same tenantID:agentID (batch queue key guarantees this). + if len(embedded) > 0 { + first := embedded[0].payload + w.classifyLinks(ctx, first.TenantID, first.AgentID, embedded) + } + + // Phase 4 — Record dedup + wikilinks. + // Dedup recorded AFTER classify so failed classify allows re-enrichment. + for _, r := range embedded { w.recordDedup(r.payload.DocID, r.payload.ContentHash) - - // Auto-link via vector similarity. - w.autoLink(ctx, r.payload.TenantID, r.payload.AgentID, r.payload.DocID) - - // Sync [[wikilinks]] from document content (text files only). w.syncWikilinks(ctx, r.payload) } @@ -210,29 +219,6 @@ func (w *enrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps } } -// autoLink finds similar documents and creates semantic links. -func (w *enrichWorker) autoLink(ctx context.Context, tenantID, agentID, docID string) { - neighbors, err := w.vault.FindSimilarDocs(ctx, tenantID, agentID, docID, enrichSimilarityLimit) - if err != nil { - slog.Warn("vault.enrich: find_similar", "doc", docID, "err", err) - return - } - - for _, n := range neighbors { - if n.Score < enrichSimilarityMin { - continue - } - link := &store.VaultLink{ - FromDocID: docID, - ToDocID: n.Document.ID, - LinkType: "semantic", - Context: fmt.Sprintf("auto-linked (score: %.2f)", n.Score), - } - if err := w.vault.CreateLink(ctx, link); err != nil { - slog.Debug("vault.enrich: create_link", "from", docID, "to", n.Document.ID, "err", err) - } - } -} // recordDedup stores a processed hash and evicts ~25% entries if over capacity. func (w *enrichWorker) recordDedup(docID, hash string) { diff --git a/tests/integration/vault_classify_test.go b/tests/integration/vault_classify_test.go new file mode 100644 index 00000000..f6d0a5a5 --- /dev/null +++ b/tests/integration/vault_classify_test.go @@ -0,0 +1,578 @@ +//go:build integration + +package integration + +import ( + "database/sql" + "errors" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/store/pg" +) + +func TestVaultClassify_DeleteDocLinksByTypes_SingleType(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create 2 documents + docA := makeVaultDoc(tid, aid, "classify/source.md", "Source Doc") + docB := makeVaultDoc(tid, aid, "classify/target.md", "Target Doc") + if err := vs.UpsertDocument(ctx, docA); err != nil { + t.Fatalf("UpsertDocument A: %v", err) + } + if err := vs.UpsertDocument(ctx, docB); err != nil { + t.Fatalf("UpsertDocument B: %v", err) + } + + // Create 3 links with different types + links := []struct { + linkType string + context string + }{ + {"reference", "cited in source"}, + {"wikilink", "mentioned in source"}, + {"depends_on", "dependency of source"}, + } + + for _, l := range links { + link := &store.VaultLink{ + FromDocID: docA.ID, + ToDocID: docB.ID, + LinkType: l.linkType, + Context: l.context, + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink %s: %v", l.linkType, err) + } + } + + // Verify: 3 links exist + outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID) + if err != nil { + t.Fatalf("GetOutLinks before delete: %v", err) + } + if len(outLinks) != 3 { + t.Errorf("expected 3 links before delete, got %d", len(outLinks)) + } + + // Delete links with type "reference" only + if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"reference"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes: %v", err) + } + + // Verify: only "reference" deleted, "wikilink" and "depends_on" remain + outLinksAfter, err := vs.GetOutLinks(ctx, tid, docA.ID) + if err != nil { + t.Fatalf("GetOutLinks after delete: %v", err) + } + if len(outLinksAfter) != 2 { + t.Errorf("expected 2 links after deleting 1, got %d", len(outLinksAfter)) + } + + // Verify remaining links are correct types + typeMap := make(map[string]bool) + for _, l := range outLinksAfter { + typeMap[l.LinkType] = true + } + if !typeMap["wikilink"] || !typeMap["depends_on"] { + t.Errorf("expected wikilink and depends_on, got types: %v", typeMap) + } + if typeMap["reference"] { + t.Errorf("reference should have been deleted") + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_MultipleTypes(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create source doc and 4 target docs (for 6 different link types + semantic) + docSource := makeVaultDoc(tid, aid, "source.md", "Source") + docs := make([]*store.VaultDocument, 0) + docs = append(docs, docSource) + + for i := 0; i < 7; i++ { + doc := makeVaultDoc(tid, aid, "classify/target-"+string(rune('a'+i))+".md", "Target "+string(rune('A'+i))) + if err := vs.UpsertDocument(ctx, doc); err != nil { + t.Fatalf("UpsertDocument target %d: %v", i, err) + } + docs = append(docs, doc) + } + docSource = docs[0] + + if err := vs.UpsertDocument(ctx, docSource); err != nil { + t.Fatalf("UpsertDocument source: %v", err) + } + + // Create links with 6 classify types + semantic + wikilink (8 total) + linkTypes := []string{ + "reference", // classify type 1 + "depends_on", // classify type 2 + "extends", // classify type 3 + "related", // classify type 4 + "supersedes", // classify type 5 + "contradicts", // classify type 6 + "semantic", // legacy auto-classified type + "wikilink", // manual link type + } + + for i, lt := range linkTypes { + link := &store.VaultLink{ + FromDocID: docSource.ID, + ToDocID: docs[i+1].ID, + LinkType: lt, + Context: "test context for " + lt, + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink %s: %v", lt, err) + } + } + + // Verify: 8 links created + outLinks, err := vs.GetOutLinks(ctx, tid, docSource.ID) + if err != nil { + t.Fatalf("GetOutLinks before delete: %v", err) + } + if len(outLinks) != 8 { + t.Errorf("expected 8 links before delete, got %d", len(outLinks)) + } + + // Delete all classify types + semantic (7 types), keep wikilink + typesToDelete := []string{ + "reference", "depends_on", "extends", + "related", "supersedes", "contradicts", + "semantic", + } + if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource.ID, typesToDelete); err != nil { + t.Fatalf("DeleteDocLinksByTypes: %v", err) + } + + // Verify: only wikilink survives + outLinksAfter, err := vs.GetOutLinks(ctx, tid, docSource.ID) + if err != nil { + t.Fatalf("GetOutLinks after delete: %v", err) + } + if len(outLinksAfter) != 1 { + t.Errorf("expected 1 link after deleting 7, got %d", len(outLinksAfter)) + } + if outLinksAfter[0].LinkType != "wikilink" { + t.Errorf("expected wikilink, got %s", outLinksAfter[0].LinkType) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_NoMatches(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create docs + docA := makeVaultDoc(tid, aid, "nomatch/a.md", "Doc A") + docB := makeVaultDoc(tid, aid, "nomatch/b.md", "Doc B") + if err := vs.UpsertDocument(ctx, docA); err != nil { + t.Fatalf("UpsertDocument A: %v", err) + } + if err := vs.UpsertDocument(ctx, docB); err != nil { + t.Fatalf("UpsertDocument B: %v", err) + } + + // Create link with type "wikilink" + link := &store.VaultLink{ + FromDocID: docA.ID, + ToDocID: docB.ID, + LinkType: "wikilink", + Context: "manual link", + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink: %v", err) + } + + // Delete non-existent type — should succeed but delete nothing + if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"reference"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes non-existent type: %v", err) + } + + // Verify wikilink still exists + outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID) + if err != nil { + t.Fatalf("GetOutLinks after no-op delete: %v", err) + } + if len(outLinks) != 1 { + t.Errorf("expected 1 link (unchanged), got %d", len(outLinks)) + } + if outLinks[0].LinkType != "wikilink" { + t.Errorf("expected wikilink, got %s", outLinks[0].LinkType) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_EmptyTypeList(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create docs and link + docA := makeVaultDoc(tid, aid, "empty/a.md", "Doc A") + docB := makeVaultDoc(tid, aid, "empty/b.md", "Doc B") + if err := vs.UpsertDocument(ctx, docA); err != nil { + t.Fatalf("UpsertDocument A: %v", err) + } + if err := vs.UpsertDocument(ctx, docB); err != nil { + t.Fatalf("UpsertDocument B: %v", err) + } + + link := &store.VaultLink{ + FromDocID: docA.ID, + ToDocID: docB.ID, + LinkType: "reference", + Context: "test", + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink: %v", err) + } + + // Delete with empty type list — should succeed but delete nothing + if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{}); err != nil { + t.Fatalf("DeleteDocLinksByTypes empty list: %v", err) + } + + // Verify link still exists + outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID) + if err != nil { + t.Fatalf("GetOutLinks after empty delete: %v", err) + } + if len(outLinks) != 1 { + t.Errorf("expected 1 link (unchanged), got %d", len(outLinks)) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_MultipleSourceDocs(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create 2 source docs and 1 target doc + docSource1 := makeVaultDoc(tid, aid, "multi/source1.md", "Source 1") + docSource2 := makeVaultDoc(tid, aid, "multi/source2.md", "Source 2") + docTarget := makeVaultDoc(tid, aid, "multi/target.md", "Target") + + for _, doc := range []*store.VaultDocument{docSource1, docSource2, docTarget} { + if err := vs.UpsertDocument(ctx, doc); err != nil { + t.Fatalf("UpsertDocument: %v", err) + } + } + + // Create links from both sources to target + linkTypes := []string{"reference", "semantic", "wikilink"} + for _, lt := range linkTypes { + link1 := &store.VaultLink{ + FromDocID: docSource1.ID, + ToDocID: docTarget.ID, + LinkType: lt, + Context: "from source1", + } + link2 := &store.VaultLink{ + FromDocID: docSource2.ID, + ToDocID: docTarget.ID, + LinkType: lt, + Context: "from source2", + } + if err := vs.CreateLink(ctx, link1); err != nil { + t.Fatalf("CreateLink source1: %v", err) + } + if err := vs.CreateLink(ctx, link2); err != nil { + t.Fatalf("CreateLink source2: %v", err) + } + } + + // Verify: both sources have 3 links each + outLinks1, err := vs.GetOutLinks(ctx, tid, docSource1.ID) + if err != nil { + t.Fatalf("GetOutLinks source1 before: %v", err) + } + outLinks2, err := vs.GetOutLinks(ctx, tid, docSource2.ID) + if err != nil { + t.Fatalf("GetOutLinks source2 before: %v", err) + } + if len(outLinks1) != 3 || len(outLinks2) != 3 { + t.Errorf("expected 3 links each source before delete, got %d and %d", len(outLinks1), len(outLinks2)) + } + + // Delete "reference" + "semantic" from source1 only + if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource1.ID, []string{"reference", "semantic"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes source1: %v", err) + } + + // Verify: source1 has only "wikilink", source2 unchanged + outLinks1After, err := vs.GetOutLinks(ctx, tid, docSource1.ID) + if err != nil { + t.Fatalf("GetOutLinks source1 after: %v", err) + } + outLinks2After, err := vs.GetOutLinks(ctx, tid, docSource2.ID) + if err != nil { + t.Fatalf("GetOutLinks source2 after: %v", err) + } + + if len(outLinks1After) != 1 { + t.Errorf("expected 1 link in source1 after delete, got %d", len(outLinks1After)) + } + if outLinks1After[0].LinkType != "wikilink" { + t.Errorf("expected wikilink in source1, got %s", outLinks1After[0].LinkType) + } + + if len(outLinks2After) != 3 { + t.Errorf("expected 3 links in source2 (unchanged), got %d", len(outLinks2After)) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_TenantIsolation(t *testing.T) { + db := testDB(t) + tenantA, tenantB, agentA, agentB := seedTwoTenants(t, db) + ctxA := tenantCtx(tenantA) + ctxB := tenantCtx(tenantB) + vs := newVaultStore(db) + + tidA := tenantA.String() + tidB := tenantB.String() + aidA := agentA.String() + aidB := agentB.String() + + // Create docs in tenant A + docA1 := makeVaultDoc(tidA, aidA, "iso/a1.md", "Tenant A Doc 1") + docA2 := makeVaultDoc(tidA, aidA, "iso/a2.md", "Tenant A Doc 2") + if err := vs.UpsertDocument(ctxA, docA1); err != nil { + t.Fatalf("UpsertDocument tenantA doc1: %v", err) + } + if err := vs.UpsertDocument(ctxA, docA2); err != nil { + t.Fatalf("UpsertDocument tenantA doc2: %v", err) + } + + // Create docs in tenant B + docB1 := makeVaultDoc(tidB, aidB, "iso/b1.md", "Tenant B Doc 1") + docB2 := makeVaultDoc(tidB, aidB, "iso/b2.md", "Tenant B Doc 2") + if err := vs.UpsertDocument(ctxB, docB1); err != nil { + t.Fatalf("UpsertDocument tenantB doc1: %v", err) + } + if err := vs.UpsertDocument(ctxB, docB2); err != nil { + t.Fatalf("UpsertDocument tenantB doc2: %v", err) + } + + // Create links in both tenants + linkA := &store.VaultLink{ + FromDocID: docA1.ID, + ToDocID: docA2.ID, + LinkType: "reference", + Context: "tenant A link", + } + linkB := &store.VaultLink{ + FromDocID: docB1.ID, + ToDocID: docB2.ID, + LinkType: "reference", + Context: "tenant B link", + } + + if err := vs.CreateLink(ctxA, linkA); err != nil { + t.Fatalf("CreateLink tenantA: %v", err) + } + if err := vs.CreateLink(ctxB, linkB); err != nil { + t.Fatalf("CreateLink tenantB: %v", err) + } + + // Verify both have links + outA, err := vs.GetOutLinks(ctxA, tidA, docA1.ID) + if err != nil { + t.Fatalf("GetOutLinks tenantA before: %v", err) + } + outB, err := vs.GetOutLinks(ctxB, tidB, docB1.ID) + if err != nil { + t.Fatalf("GetOutLinks tenantB before: %v", err) + } + if len(outA) != 1 || len(outB) != 1 { + t.Errorf("expected 1 link each before delete, got %d and %d", len(outA), len(outB)) + } + + // Delete from tenant A only — should NOT affect tenant B + if err := vs.DeleteDocLinksByTypes(ctxA, tidA, docA1.ID, []string{"reference"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes tenantA: %v", err) + } + + // Verify: tenantA link deleted, tenantB link unchanged + outAAfter, err := vs.GetOutLinks(ctxA, tidA, docA1.ID) + if err != nil { + t.Fatalf("GetOutLinks tenantA after: %v", err) + } + outBAfter, err := vs.GetOutLinks(ctxB, tidB, docB1.ID) + if err != nil { + t.Fatalf("GetOutLinks tenantB after: %v", err) + } + + if len(outAAfter) != 0 { + t.Errorf("expected 0 links in tenantA after delete, got %d", len(outAAfter)) + } + if len(outBAfter) != 1 { + t.Errorf("expected 1 link in tenantB (unchanged), got %d", len(outBAfter)) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_NonExistentDoc(t *testing.T) { + db := testDB(t) + tenantID, _ := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + + // Try to delete links from non-existent doc — should succeed (no-op) + fakeDocID := "fake-doc-id-12345" + if err := vs.DeleteDocLinksByTypes(ctx, tid, fakeDocID, []string{"reference"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes on non-existent doc: %v", err) + } + + // GetOutLinks for non-existent doc should return empty or error gracefully + outLinks, err := vs.GetOutLinks(ctx, tid, fakeDocID) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + // Some implementations may error on non-existent doc, others return empty list + // Both are acceptable for a non-existent doc + } + if outLinks != nil && len(outLinks) != 0 { + t.Errorf("expected empty or error for non-existent doc, got %d links", len(outLinks)) + } +} + +func TestVaultClassify_DeleteAllClassifyTypesAndSemantic(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create documents + docSource := makeVaultDoc(tid, aid, "classify/source.md", "Source Doc") + docTarget1 := makeVaultDoc(tid, aid, "classify/target1.md", "Target 1") + docTarget2 := makeVaultDoc(tid, aid, "classify/target2.md", "Target 2") + + for _, doc := range []*store.VaultDocument{docSource, docTarget1, docTarget2} { + if err := vs.UpsertDocument(ctx, doc); err != nil { + t.Fatalf("UpsertDocument: %v", err) + } + } + + // Create links with mix of classify types, semantic, and manual types + linkConfigs := []struct { + toDocID string + linkType string + }{ + {docTarget1.ID, "reference"}, + {docTarget1.ID, "semantic"}, + {docTarget2.ID, "depends_on"}, + {docTarget2.ID, "wikilink"}, + } + + for _, cfg := range linkConfigs { + link := &store.VaultLink{ + FromDocID: docSource.ID, + ToDocID: cfg.toDocID, + LinkType: cfg.linkType, + Context: "test context", + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink %s: %v", cfg.linkType, err) + } + } + + // Verify: 4 links created + outLinks, err := vs.GetOutLinks(ctx, tid, docSource.ID) + if err != nil { + t.Fatalf("GetOutLinks before delete: %v", err) + } + if len(outLinks) != 4 { + t.Errorf("expected 4 links before delete, got %d", len(outLinks)) + } + + // Delete all classify types (reference, depends_on) + semantic, keep wikilink + if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource.ID, []string{ + "reference", "depends_on", "extends", "related", "supersedes", "contradicts", "semantic", + }); err != nil { + t.Fatalf("DeleteDocLinksByTypes: %v", err) + } + + // Verify: only wikilink remains + outLinksAfter, err := vs.GetOutLinks(ctx, tid, docSource.ID) + if err != nil { + t.Fatalf("GetOutLinks after delete: %v", err) + } + if len(outLinksAfter) != 1 { + t.Errorf("expected 1 link after cleanup, got %d", len(outLinksAfter)) + } + if outLinksAfter[0].LinkType != "wikilink" { + t.Errorf("expected wikilink survivor, got %s", outLinksAfter[0].LinkType) + } +} + +func TestVaultClassify_DeleteDocLinksByTypes_CasePreservation(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vs := newVaultStore(db) + + tid := tenantID.String() + aid := agentID.String() + + // Create documents + docA := makeVaultDoc(tid, aid, "case/a.md", "Doc A") + docB := makeVaultDoc(tid, aid, "case/b.md", "Doc B") + if err := vs.UpsertDocument(ctx, docA); err != nil { + t.Fatalf("UpsertDocument A: %v", err) + } + if err := vs.UpsertDocument(ctx, docB); err != nil { + t.Fatalf("UpsertDocument B: %v", err) + } + + // Create links with different case sensitivity + link := &store.VaultLink{ + FromDocID: docA.ID, + ToDocID: docB.ID, + LinkType: "depends_on", // lowercase with underscore + Context: "test", + } + if err := vs.CreateLink(ctx, link); err != nil { + t.Fatalf("CreateLink: %v", err) + } + + // Try deleting with exact case match + if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"depends_on"}); err != nil { + t.Fatalf("DeleteDocLinksByTypes: %v", err) + } + + // Verify link is deleted + outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID) + if err != nil { + t.Fatalf("GetOutLinks: %v", err) + } + if len(outLinks) != 0 { + t.Errorf("expected 0 links after exact case match delete, got %d", len(outLinks)) + } +}