From c13c4ec371cc5194f86923c05a80446ca18fe00f Mon Sep 17 00:00:00 2001 From: avfirsov Date: Tue, 23 Jun 2026 12:50:13 +0300 Subject: [PATCH] =?UTF-8?q?feat(graph,mcp):=20annotate=5Fnodes=20=E2=80=94?= =?UTF-8?q?=20merge=20external=20metadata=20onto=20graph=20nodes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a sanctioned, additive write-back path so a downstream tool (an LLM, an analysis pass, an editor extension) can enrich existing graph nodes with its own metadata — without re-indexing and without racing the sharded in-memory store. - graph.Store.MergeNodeMeta(id, kv) (changed, found): the only sanctioned way for a Store holder to mutate Node.Meta. Additive + idempotent (deep-equal delta via metaDelta), shard-locked, and never touches structural fields (id/kind/name/path/lines). Implemented for both the in-memory and SQLite backends and covered by a shared store-conformance case so they behave identically. - mcp annotate_nodes tool (also served at POST /v1/tools/annotate_nodes through the shared registry): merges a free-form per-node `meta` map under a caller-chosen `namespace` (default "ext") so annotations can never shadow indexer-owned keys, and optionally adds idempotent semantically_related edges between node pairs. Tests: pure metaDelta unit, in-memory + SQLite conformance, and MCP handler round-trip / idempotency / namespace / edge / bad-input plus a registration guard. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01XcQY8fFFyDQidwVC8dKHWh --- internal/graph/meta_merge.go | 89 ++++++++++++ internal/graph/meta_merge_test.go | 199 ++++++++++++++++++++++++++ internal/graph/store.go | 32 +++++ internal/graph/store_sqlite/store.go | 63 ++++++++ internal/graph/storetest/storetest.go | 66 +++++++++ internal/mcp/server.go | 1 + internal/mcp/tools_annotate.go | 198 +++++++++++++++++++++++++ internal/mcp/tools_annotate_test.go | 187 ++++++++++++++++++++++++ 8 files changed, 835 insertions(+) create mode 100644 internal/graph/meta_merge.go create mode 100644 internal/graph/meta_merge_test.go create mode 100644 internal/mcp/tools_annotate.go create mode 100644 internal/mcp/tools_annotate_test.go diff --git a/internal/graph/meta_merge.go b/internal/graph/meta_merge.go new file mode 100644 index 00000000..f4a75fb6 --- /dev/null +++ b/internal/graph/meta_merge.go @@ -0,0 +1,89 @@ +package graph + +import "reflect" + +// metaDelta returns the subset of kv whose value differs from (or is +// absent in) existing — i.e. the keys a merge would actually change. +// +// It is a pure calculation: no locks, no I/O, no mutation of either +// argument. Equality is reflect.DeepEqual so the comparison is correct +// for the value shapes that ride on Node.Meta — scalars (string, the +// JSON-decoded float64 for numbers), []string / []any tag slices, and +// nested map[string]any blobs — where a == comparison would either +// panic (slices/maps are not comparable) or compare identity rather +// than contents. +// +// A nil existing is treated as "no keys present", so every entry in kv +// is returned. An empty/nil kv returns nil (nothing to merge). The +// returned map is freshly allocated and shares no structure with kv's +// values beyond the value pointers themselves (callers store them as-is +// — these are caller-owned, JSON-decoded values). +// +// Idempotency falls straight out of this: once a key has been written, +// a second metaDelta over the now-updated map omits it, so the second +// MergeNodeMeta reports changed=false. +func metaDelta(existing, kv map[string]any) map[string]any { + if len(kv) == 0 { + return nil + } + var delta map[string]any + for k, v := range kv { + if existing != nil { + if cur, ok := existing[k]; ok && reflect.DeepEqual(cur, v) { + // Key already present with an equal value — nothing to do. + continue + } + } + if delta == nil { + delta = make(map[string]any, len(kv)) + } + delta[k] = v + } + return delta +} + +// MergeNodeMeta is the in-memory *Graph implementation of the +// Store.MergeNodeMeta contract: additive, idempotent, shard-locked +// merge of kv into the target node's Meta. See the Store interface doc +// for the full (changed, found) semantics. +// +// Locking: the node lives in exactly one shard (shardFor(id)), so a +// single shard write lock is sufficient and necessary — necessary +// because GetNode reads, and a concurrent AddNode/AddEdge writes, the +// same sharded maps; sufficient because the merge touches only this one +// node's Meta map and no cross-shard index. We reuse lockTwoWrite(id, +// id), which collapses to a single Lock on that shard (its a==b branch), +// rather than hand-rolling shardFor(id).mu.Lock() so this method stays +// in lock-step with the one sanctioned shard-locking helper the rest of +// the mutators use. +// +// The compare-before-write (metaDelta) runs inside the lock: it must +// observe the same Meta it is about to mutate, and computing it outside +// the lock would race a concurrent merge to the same node. +func (g *Graph) MergeNodeMeta(id string, kv map[string]any) (changed bool, found bool) { + unlock := g.lockTwoWrite(id, id) + defer unlock() + + s := g.shardFor(id) + n := s.nodes[id] + if n == nil { + // Unknown id — recorded as a not-found skip by the caller; the + // batch continues. No structural state was touched. + return false, false + } + + delta := metaDelta(n.Meta, kv) + if len(delta) == 0 { + // Found, but every provided key already equals the stored value: + // the merge is a no-op, so the call is idempotent. + return false, true + } + + if n.Meta == nil { + n.Meta = make(map[string]any, len(delta)) + } + for k, v := range delta { + n.Meta[k] = v + } + return true, true +} diff --git a/internal/graph/meta_merge_test.go b/internal/graph/meta_merge_test.go new file mode 100644 index 00000000..bad31c55 --- /dev/null +++ b/internal/graph/meta_merge_test.go @@ -0,0 +1,199 @@ +package graph + +import ( + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestMetaDelta_OnlyDifferingKeys proves the pure delta calculation: +// metaDelta returns exactly the keys whose value is absent or differs +// from the existing map, across the value shapes that ride on Node.Meta +// (string, []any tag slice, float64, nested map) using reflect.DeepEqual +// so slices/maps compare by contents, not identity. (AC1, pure half.) +func TestMetaDelta_OnlyDifferingKeys(t *testing.T) { + existing := map[string]any{ + "ext_summary": "old summary", + "ext_tags": []any{"a", "b"}, + "ext_complexity": 0.5, + "structural": "untouched", + } + kv := map[string]any{ + "ext_summary": "new summary", // changed + "ext_tags": []any{"a", "b"}, // equal (deep) -> omitted + "ext_complexity": 0.5, // equal -> omitted + "ext_domain": "billing", // new key + "nested": map[string]any{"x": 1}, // new nested map + } + + delta := metaDelta(existing, kv) + + // Telemetry before asserts (LDD spirit). + t.Logf("[delta] keys=%v", keysOf(delta)) + + require.Len(t, delta, 3, "only changed/new keys should appear") + assert.Equal(t, "new summary", delta["ext_summary"]) + assert.Equal(t, "billing", delta["ext_domain"]) + assert.Equal(t, map[string]any{"x": 1}, delta["nested"]) + _, hasTags := delta["ext_tags"] + assert.False(t, hasTags, "deep-equal tag slice must be omitted") + _, hasComplexity := delta["ext_complexity"] + assert.False(t, hasComplexity, "equal complexity must be omitted") +} + +// TestMetaDelta_NilAndEmpty covers the boundary inputs: a nil existing +// map returns every kv entry; an empty/nil kv returns nil. +func TestMetaDelta_NilAndEmpty(t *testing.T) { + full := metaDelta(nil, map[string]any{"ext_summary": "s", "ext_domain": "d"}) + require.Len(t, full, 2, "nil existing -> every kv key is a delta") + + assert.Nil(t, metaDelta(map[string]any{"k": "v"}, nil), "nil kv -> nil delta") + assert.Nil(t, metaDelta(nil, map[string]any{}), "empty kv -> nil delta") +} + +// TestMergeNodeMeta_MergeOverwriteIdempotent exercises the full +// (changed, found) contract on the in-memory store: new keys merge, +// changed keys overwrite, and a re-apply of identical input is a no-op +// (idempotent). (AC1.) +func TestMergeNodeMeta_MergeOverwriteIdempotent(t *testing.T) { + g := New() + g.AddNode(&Node{ID: "repo/a.go::Foo", Name: "Foo", Kind: KindFunction, FilePath: "repo/a.go"}) + + // First merge: two new keys. + changed, found := g.MergeNodeMeta("repo/a.go::Foo", map[string]any{ + "ext_summary": "computes foo", + "ext_domain": "core", + }) + t.Logf("[merge#1] changed=%v found=%v", changed, found) + require.True(t, found) + require.True(t, changed, "first merge writes new keys") + assert.Equal(t, "computes foo", g.GetNode("repo/a.go::Foo").Meta["ext_summary"]) + + // Idempotent re-apply: identical input -> no change. + changed, found = g.MergeNodeMeta("repo/a.go::Foo", map[string]any{ + "ext_summary": "computes foo", + "ext_domain": "core", + }) + t.Logf("[merge#2 idempotent] changed=%v found=%v", changed, found) + require.True(t, found) + assert.False(t, changed, "re-applying identical input must be a no-op") + + // Overwrite: one key changes -> changed=true, only that key updated. + changed, found = g.MergeNodeMeta("repo/a.go::Foo", map[string]any{ + "ext_summary": "computes foo, revised", + "ext_domain": "core", // unchanged + }) + t.Logf("[merge#3 overwrite] changed=%v found=%v", changed, found) + require.True(t, found) + require.True(t, changed, "a differing value flips changed") + assert.Equal(t, "computes foo, revised", g.GetNode("repo/a.go::Foo").Meta["ext_summary"]) + assert.Equal(t, "core", g.GetNode("repo/a.go::Foo").Meta["ext_domain"]) +} + +// TestMergeNodeMeta_UnknownIDNoPanic proves an unknown id reports +// found=false (changed=false) and never panics. (AC1, MUST NOT: unknown +// ids must not crash the batch.) +func TestMergeNodeMeta_UnknownIDNoPanic(t *testing.T) { + g := New() + changed, found := g.MergeNodeMeta("does/not/exist::Nope", map[string]any{"ext_summary": "x"}) + t.Logf("[unknown] changed=%v found=%v", changed, found) + assert.False(t, found, "unknown id -> found=false") + assert.False(t, changed, "unknown id -> changed=false") +} + +// TestMergeNodeMeta_LazyInitsNilMeta proves a node created without a +// Meta map gets one lazily on first merge — no nil-map write panic. +func TestMergeNodeMeta_LazyInitsNilMeta(t *testing.T) { + g := New() + g.AddNode(&Node{ID: "repo/b.go::Bar", Name: "Bar", Kind: KindFunction, FilePath: "repo/b.go"}) + require.Nil(t, g.GetNode("repo/b.go::Bar").Meta, "precondition: nil Meta") + + changed, found := g.MergeNodeMeta("repo/b.go::Bar", map[string]any{"ext_summary": "lazy"}) + require.True(t, found) + require.True(t, changed) + assert.Equal(t, "lazy", g.GetNode("repo/b.go::Bar").Meta["ext_summary"]) +} + +// TestMergeNodeMeta_StructuralUntouched proves the merge never alters +// structural node fields. (AC3 / MUST NOT.) +func TestMergeNodeMeta_StructuralUntouched(t *testing.T) { + g := New() + g.AddNode(&Node{ID: "repo/c.go::Baz", Name: "Baz", Kind: KindFunction, FilePath: "repo/c.go", StartLine: 10, EndLine: 20}) + + g.MergeNodeMeta("repo/c.go::Baz", map[string]any{"ext_summary": "s", "ext_tags": []any{"t"}}) + + n := g.GetNode("repo/c.go::Baz") + assert.Equal(t, "Baz", n.Name) + assert.Equal(t, KindFunction, n.Kind) + assert.Equal(t, "repo/c.go", n.FilePath) + assert.Equal(t, 10, n.StartLine) + assert.Equal(t, 20, n.EndLine) + assert.Equal(t, "s", n.Meta["ext_summary"]) +} + +// TestMergeNodeMeta_ConcurrentNoRace fans out many MergeNodeMeta and +// AddEdge goroutines against a shared graph. Run under `-race` it proves +// the merge takes the shard write lock correctly: no concurrent-map +// panic, no data race. This is the first mutating tool, so this is the +// load-bearing safety test. (AC2.) +func TestMergeNodeMeta_ConcurrentNoRace(t *testing.T) { + g := New() + const nNodes = 32 + for i := 0; i < nNodes; i++ { + g.AddNode(&Node{ + ID: fmt.Sprintf("repo/x.go::N%d", i), + Name: fmt.Sprintf("N%d", i), + Kind: KindFunction, + FilePath: "repo/x.go", + }) + } + + const workers = 16 + const iters = 200 + var wg sync.WaitGroup + wg.Add(workers) + for w := 0; w < workers; w++ { + go func(w int) { + defer wg.Done() + for it := 0; it < iters; it++ { + id := fmt.Sprintf("repo/x.go::N%d", (w+it)%nNodes) + // Concurrent Meta merge (shard-locked) ... + g.MergeNodeMeta(id, map[string]any{ + "ext_summary": fmt.Sprintf("w%d-it%d", w, it), + "ext_tags": []any{"t", fmt.Sprintf("%d", it)}, + }) + // ... interleaved with idempotent edge adds (also + // shard-locked) to maximise cross-shard lock contention. + other := fmt.Sprintf("repo/x.go::N%d", (w+it+1)%nNodes) + g.AddEdge(&Edge{ + From: id, + To: other, + Kind: EdgeSemanticallyRelated, + Origin: "ext_annotated", + }) + } + }(w) + } + wg.Wait() + + t.Logf("[concurrency] nodes=%d edges=%d", g.NodeCount(), g.EdgeCount()) + // Every node should carry a ext_summary from the last writer that + // touched it; the exact value is non-deterministic, presence is not. + for i := 0; i < nNodes; i++ { + n := g.GetNode(fmt.Sprintf("repo/x.go::N%d", i)) + require.NotNil(t, n) + assert.Contains(t, n.Meta, "ext_summary") + } +} + +// keysOf is a tiny test helper for stable delta-key logging. +func keysOf(m map[string]any) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} diff --git a/internal/graph/store.go b/internal/graph/store.go index e4878d97..f2bd15fe 100644 --- a/internal/graph/store.go +++ b/internal/graph/store.go @@ -60,6 +60,38 @@ type Store interface { AddNode(n *Node) AddBatch(nodes []*Node, edges []*Edge) AddEdge(e *Edge) + // MergeNodeMeta merges the key/value pairs in kv into the Meta map + // of the node identified by id, under the node's shard write lock. + // It is the ONLY sanctioned way for a caller holding a Store (rather + // than a concrete *Graph) to mutate Node.Meta: the shard lock the + // in-memory backend takes internally is private, so a handler that + // reached into GetNode(id).Meta directly would race the sharded map + // and panic on a live daemon. + // + // Semantics: + // + // - found is false (and changed is false) when no node carries id — + // callers batching many ids treat a miss as a recorded skip, not + // a failure. + // - changed is false when every key in kv already maps to an equal + // value (deep-equal compare via metaDelta) — the merge is a no-op + // and the call is idempotent. Re-applying the same kv twice yields + // changed=false on the second call. + // - changed is true when at least one key was added or its value + // differed; only the differing keys are written, and a nil Meta is + // lazily initialised before the first write. + // + // The method is deliberately scoped to additive Meta merges — it + // never deletes keys and never touches structural fields (ID, Kind, + // Name, FilePath, line ranges). Callers writing external/agent-produced + // annotations namespace their keys (e.g. ext_summary, ext_tags, …) so a + // merge can never shadow an indexer-owned Meta key by accident. + // + // Durability: the merge mutates in-memory state for the daemon + // session and rides the gob+gzip shutdown snapshot; there is no + // per-call persist. Cross-restart durability of annotations is an + // out-of-scope follow-up for this v1. + MergeNodeMeta(id string, kv map[string]any) (changed bool, found bool) SetEdgeProvenance(e *Edge, newOrigin string) bool ReindexEdge(e *Edge, oldTo string) // Batched siblings of the per-edge mutators. Same semantics, but diff --git a/internal/graph/store_sqlite/store.go b/internal/graph/store_sqlite/store.go index 3ebdeb8b..0481c70c 100644 --- a/internal/graph/store_sqlite/store.go +++ b/internal/graph/store_sqlite/store.go @@ -26,6 +26,7 @@ import ( "errors" "fmt" "iter" + "reflect" "runtime" "strings" "sync" @@ -651,6 +652,68 @@ func (s *Store) insertNodeLocked(stmt *sql.Stmt, n *graph.Node) error { return err } +// MergeNodeMeta is the on-disk implementation of the +// Store.MergeNodeMeta contract: an additive, idempotent merge of kv +// into the target node's Meta. See the graph.Store interface doc for +// the full (changed, found) semantics. +// +// The whole read-modify-write runs under writeMu: the node is read +// back (decoding its gob Meta blob), the differing keys are computed, +// and — only if at least one key changed — the node is re-persisted via +// the INSERT OR REPLACE node statement. Holding writeMu across the read +// is what makes the merge atomic against a concurrent MergeNodeMeta / +// AddNode to the same id; a read outside the lock could lose an update. +// +// Equality uses reflect.DeepEqual for the same reason the in-memory +// metaDelta does: Meta values are JSON/gob-decoded scalars, slices, and +// nested maps where == is either a panic or an identity compare. +func (s *Store) MergeNodeMeta(id string, kv map[string]any) (changed bool, found bool) { + if len(kv) == 0 { + // Nothing to merge — report found honestly without a write. + if n := s.GetNode(id); n == nil { + return false, false + } + return false, true + } + + s.writeMu.Lock() + defer s.writeMu.Unlock() + + n := s.GetNode(id) + if n == nil { + return false, false + } + + // Compute the differing keys against the freshly-read Meta. + delta := make(map[string]any, len(kv)) + for k, v := range kv { + if n.Meta != nil { + if cur, ok := n.Meta[k]; ok && reflect.DeepEqual(cur, v) { + continue + } + } + delta[k] = v + } + if len(delta) == 0 { + // Found, but already up to date — idempotent no-op. + return false, true + } + + if n.Meta == nil { + n.Meta = make(map[string]any, len(delta)) + } + for k, v := range delta { + n.Meta[k] = v + } + if err := s.insertNodeLocked(s.stmtInsertNode, n); err != nil { + // graph.Store.MergeNodeMeta has no error channel; mirror AddNode + // and panic only on a clearly catastrophic failure. + panicOnFatal(err) + return false, true + } + return true, true +} + // AddEdge inserts an edge. Idempotent on the logical edge key (from, // to, kind, file_path, line) -- a second AddEdge with the same key is // a no-op (INSERT OR IGNORE), matching the in-memory store's "stored diff --git a/internal/graph/storetest/storetest.go b/internal/graph/storetest/storetest.go index 163d743e..58e016ae 100644 --- a/internal/graph/storetest/storetest.go +++ b/internal/graph/storetest/storetest.go @@ -63,6 +63,7 @@ func RunConformance(t *testing.T, factory Factory) { t.Run("RepoMemoryEstimate", func(t *testing.T) { testRepoMemoryEstimate(t, factory) }) t.Run("AllRepoMemoryEstimates", func(t *testing.T) { testAllRepoMemoryEstimates(t, factory) }) t.Run("MetaPreserved", func(t *testing.T) { testMetaPreserved(t, factory) }) + t.Run("MergeNodeMeta", func(t *testing.T) { testMergeNodeMeta(t, factory) }) t.Run("EmptyStore", func(t *testing.T) { testEmptyStore(t, factory) }) t.Run("EdgesByKind", func(t *testing.T) { testEdgesByKind(t, factory) }) t.Run("NodesByKind", func(t *testing.T) { testNodesByKind(t, factory) }) @@ -105,6 +106,71 @@ func RunConformance(t *testing.T, factory Factory) { t.Run("ContractBridgeRoundTrip", func(t *testing.T) { testContractBridgeRoundTrip(t, factory) }) } +// testMergeNodeMeta verifies the Store.MergeNodeMeta contract across +// every backend: additive merge, deep-equal idempotency, found +// semantics for an unknown id, lazy init of a nil Meta, and that +// structural fields stay untouched. The in-memory store has its own +// finer-grained tests in internal/graph; running the same contract here +// proves the SQLite backend honours it identically (read-modify-write +// of the gob Meta blob under writeMu). +func testMergeNodeMeta(t *testing.T, factory Factory) { + t.Helper() + s := factory(t) + s.AddNode(mkNode("a.go::Foo", "Foo", "a.go", graph.KindFunction)) + + // First merge: two new keys -> changed=true, found=true. + changed, found := s.MergeNodeMeta("a.go::Foo", map[string]any{ + "ext_summary": "computes foo", + "ext_domain": "core", + }) + if !found || !changed { + t.Fatalf("first merge: changed=%v found=%v, want true/true", changed, found) + } + got := s.GetNode("a.go::Foo") + if got == nil || got.Meta["ext_summary"] != "computes foo" || got.Meta["ext_domain"] != "core" { + t.Fatalf("merged Meta not persisted: %+v", got) + } + + // Idempotent re-apply: identical input -> changed=false. + changed, found = s.MergeNodeMeta("a.go::Foo", map[string]any{ + "ext_summary": "computes foo", + "ext_domain": "core", + }) + if !found || changed { + t.Fatalf("idempotent merge: changed=%v found=%v, want false/true", changed, found) + } + + // Overwrite one key -> changed=true. + changed, found = s.MergeNodeMeta("a.go::Foo", map[string]any{"ext_summary": "revised"}) + if !found || !changed { + t.Fatalf("overwrite merge: changed=%v found=%v, want true/true", changed, found) + } + got = s.GetNode("a.go::Foo") + if got.Meta["ext_summary"] != "revised" || got.Meta["ext_domain"] != "core" { + t.Fatalf("overwrite Meta wrong: %+v", got.Meta) + } + // Structural fields untouched. + if got.Name != "Foo" || got.Kind != graph.KindFunction || got.FilePath != "a.go" { + t.Fatalf("structural fields mutated: %+v", got) + } + + // Unknown id -> found=false, no panic. + changed, found = s.MergeNodeMeta("a.go::Ghost", map[string]any{"ext_summary": "x"}) + if found || changed { + t.Fatalf("unknown id: changed=%v found=%v, want false/false", changed, found) + } + + // Lazy-init: a node added with nil Meta gets one on first merge. + s.AddNode(mkNode("b.go::Bar", "Bar", "b.go", graph.KindFunction)) + changed, found = s.MergeNodeMeta("b.go::Bar", map[string]any{"ext_summary": "lazy"}) + if !found || !changed { + t.Fatalf("lazy-init merge: changed=%v found=%v, want true/true", changed, found) + } + if s.GetNode("b.go::Bar").Meta["ext_summary"] != "lazy" { + t.Fatalf("lazy-init Meta not persisted") + } +} + // -- fixture helpers --------------------------------------------------- func mkNode(id, name, file string, kind graph.NodeKind) *graph.Node { diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 1455b38e..39ac2ce4 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -1282,6 +1282,7 @@ func NewServer(engine *query.Engine, g graph.Store, idx *indexer.Indexer, watche s.registerGraphCompletionTool() s.registerWikiTools() s.registerExportTools() + s.registerAnnotateTools() s.registerAuditTool() s.registerWalkGraphTool() s.registerContextClosureTool() diff --git a/internal/mcp/tools_annotate.go b/internal/mcp/tools_annotate.go new file mode 100644 index 00000000..61426ebc --- /dev/null +++ b/internal/mcp/tools_annotate.go @@ -0,0 +1,198 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/mark3labs/mcp-go/mcp" + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/graph" +) + +// defaultAnnotateNamespace is the key prefix applied to merged metadata +// when the caller does not pass an explicit `namespace`. "ext" reads as +// "external" — metadata produced outside the indexer (an agent, an LLM, +// a downstream tool) — and keeps annotation keys clearly distinct from +// the indexer-owned Meta keys. +const defaultAnnotateNamespace = "ext" + +// annotateInput is one entry in the annotate_nodes `annotations` array: +// a graph node id plus a free-form metadata map to merge onto it. The +// map is intentionally untyped — callers write whatever +// external/agent-produced fields they like (summary, tags, complexity, +// domain, …); only the keys actually present are merged, so an omitted +// field never clobbers a prior annotation. +type annotateInput struct { + ID string `json:"id"` + Meta map[string]any `json:"meta,omitempty"` +} + +// namespacedKV builds the Meta key/value map that is merged into a node, +// prefixing every key with "_". The prefix is what preserves +// the no-clobber invariant: indexer-owned Meta keys are unprefixed, so a +// namespaced annotation can never shadow one by accident. A key that +// already carries the prefix is left as-is, so a caller may pass either +// the bare ("summary") or fully-qualified ("ext_summary") form. +// +// An entry with an empty Meta map yields nil — MergeNodeMeta then reports +// the node found but unchanged, which the handler counts as "unchanged". +// +// Values are stored exactly as JSON decoded them (string, float64 for +// numbers, []any for arrays, map[string]any for objects). That is the +// shape MergeNodeMeta's reflect.DeepEqual idempotency check expects, so +// re-merging an identical annotation is a stable no-op. +func (a annotateInput) namespacedKV(namespace string) map[string]any { + if len(a.Meta) == 0 { + return nil + } + prefix := namespace + "_" + kv := make(map[string]any, len(a.Meta)) + for k, v := range a.Meta { + if k == "" { + continue + } + key := k + if !strings.HasPrefix(key, prefix) { + key = prefix + key + } + kv[key] = v + } + return kv +} + +// registerAnnotateTools wires the annotate_nodes write-back tool onto the +// MCP tool surface. Registering it here is sufficient to also serve +// POST /v1/tools/annotate_nodes: the daemon's HTTP handler dispatches +// against this same MCP registry, so no separate HTTP code is needed. +func (s *Server) registerAnnotateTools() { + s.addTool( + mcp.NewTool("annotate_nodes", + mcp.WithDescription("Merge external / agent-produced metadata back into existing graph nodes by id. Each annotation's `meta` keys are namespaced (default `ext_`) before merging, so they can never overwrite indexer-owned Meta — the merge is idempotent and additive and never mutates structural data (id/kind/name/path/lines). Optionally adds semantically_related edges between node pairs. Returns {annotated, unchanged, not_found, edges_added}."), + mcp.WithString("annotations", mcp.Required(), mcp.Description(`JSON array of per-node annotations: [{"id":"","meta":{"summary":"...","tags":["..."],"complexity":0.7,"domain":"..."}}]. The meta object is free-form; only present keys are merged, and each key is namespaced before the merge.`)), + mcp.WithString("namespace", mcp.Description(`Key namespace prefix for merged meta (default "ext"). Every meta key is stored as "_" so annotations never shadow indexer-owned keys. A key already carrying the prefix is left as-is.`)), + mcp.WithString("add_related", mcp.Description(`Optional JSON array of semantically_related edge pairs: [["idA","idB",0.8]]. The third element (score, 0..1) is optional and defaults to 0.5.`)), + ), + s.handleAnnotateNodes, + ) +} + +// handleAnnotateNodes is the Action-layer controller for annotate_nodes. +// It parses the annotations + add_related JSON arguments, merges each +// node's namespaced metadata via the shard-locked graph.Store.MergeNodeMeta +// (the only sanctioned Meta-mutation path), adds any requested +// semantically_related edges via the idempotent AddEdge, and returns the +// {annotated, unchanged, not_found, edges_added} summary. +// +// Failure handling is deliberately non-fatal per item: a node id that +// is not in the graph is recorded in not_found and the batch continues +// — one bad id never fails the whole call. Structural data is never +// touched: the only writes are additive namespaced Meta keys and +// semantically_related edges. +func (s *Server) handleAnnotateNodes(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + g := s.graph + if g == nil { + return mcp.NewToolResultError("annotate_nodes: graph is not initialised"), nil + } + args := req.GetArguments() + + // --- Resolve the key namespace --------------------------------- + namespace := stringArg(args, "namespace") + if namespace == "" { + namespace = defaultAnnotateNamespace + } + + // --- Parse the annotations array (required) --------------------- + rawAnnotations := stringArg(args, "annotations") + if rawAnnotations == "" { + return mcp.NewToolResultError("annotate_nodes: 'annotations' is required (a JSON array of {id, meta})"), nil + } + var annotations []annotateInput + if err := json.Unmarshal([]byte(rawAnnotations), &annotations); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("annotate_nodes: 'annotations' is not valid JSON: %v", err)), nil + } + + // --- Merge each node's namespaced metadata ---------------------- + annotated := 0 + unchanged := 0 + notFound := make([]string, 0) + for _, a := range annotations { + if a.ID == "" { + // An entry with no id can't address a node — treat as a + // not-found skip (recorded under the empty id) rather than + // failing the batch. + notFound = append(notFound, "") + continue + } + changed, found := g.MergeNodeMeta(a.ID, a.namespacedKV(namespace)) + switch { + case !found: + notFound = append(notFound, a.ID) + case changed: + annotated++ + default: + unchanged++ + } + } + + // --- Add optional semantically_related edges -------------------- + edgesAdded := 0 + if rawRelated := stringArg(args, "add_related"); rawRelated != "" { + var pairs [][]any + if err := json.Unmarshal([]byte(rawRelated), &pairs); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("annotate_nodes: 'add_related' is not valid JSON: %v", err)), nil + } + for _, p := range pairs { + if len(p) < 2 { + continue + } + a, aok := p[0].(string) + b, bok := p[1].(string) + if !aok || !bok || a == "" || b == "" { + continue + } + score := scoreOrDefault(p) + g.AddEdge(&graph.Edge{ + From: a, + To: b, + Kind: graph.EdgeSemanticallyRelated, + Confidence: score, + Origin: namespace + "_annotated", + Meta: map[string]any{"similarity": score}, + }) + edgesAdded++ + } + } + + if s.logger != nil { + s.logger.Info("annotate_nodes", + zap.String("namespace", namespace), + zap.Int("annotated", annotated), + zap.Int("unchanged", unchanged), + zap.Int("not_found", len(notFound)), + zap.Int("edges_added", edgesAdded), + ) + } + + return s.respondJSONOrTOON(ctx, req, map[string]any{ + "annotated": annotated, + "unchanged": unchanged, + "not_found": notFound, + "edges_added": edgesAdded, + }) +} + +// scoreOrDefault extracts the optional similarity score (third element) +// from an add_related pair. JSON numbers decode as float64; a missing or +// non-numeric third element falls back to 0.5 — the same default the +// graph-diffusion pass uses for an unscored semantically_related edge. +func scoreOrDefault(pair []any) float64 { + if len(pair) >= 3 { + if f, ok := pair[2].(float64); ok { + return f + } + } + return 0.5 +} diff --git a/internal/mcp/tools_annotate_test.go b/internal/mcp/tools_annotate_test.go new file mode 100644 index 00000000..b702f668 --- /dev/null +++ b/internal/mcp/tools_annotate_test.go @@ -0,0 +1,187 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" +) + +// newAnnotateTestServer builds a minimal Server over a small in-memory +// graph with the annotate tool registered, reusing the same field set +// the notes/handler tests use. The two seed nodes are the annotation +// targets for the handler tests below. +func newAnnotateTestServer(t *testing.T) *Server { + t.Helper() + g := graph.New() + g.AddNode(&graph.Node{ID: "pkg/foo.go::Bar", Name: "Bar", Kind: graph.KindFunction, FilePath: "pkg/foo.go", StartLine: 1, EndLine: 9}) + g.AddNode(&graph.Node{ID: "pkg/foo.go::Baz", Name: "Baz", Kind: graph.KindMethod, FilePath: "pkg/foo.go", StartLine: 11, EndLine: 20}) + + s := &Server{ + graph: g, + session: newSessionState(), + tokenStats: &tokenStats{}, + symHistory: &symbolHistory{entries: make(map[string][]SymbolModification)}, + sessions: newSessionMap(), + toolScopes: newScopeRegistry(), + } + return s +} + +// TestHandleAnnotateNodes_RoundTripAndIdempotent is the core handler +// test: annotate two nodes inline, assert the JSON summary, assert +// GetNode reflects the merged namespaced keys with structural fields +// untouched, then re-run the identical batch and assert it reports +// everything unchanged (idempotent). +func TestHandleAnnotateNodes_RoundTripAndIdempotent(t *testing.T) { + s := newAnnotateTestServer(t) + + annotations := `[ + {"id":"pkg/foo.go::Bar","meta":{"summary":"parses the bar","tags":["parser","io"],"complexity":0.7,"domain":"ingest"}}, + {"id":"pkg/foo.go::Baz","meta":{"summary":"formats the baz"}}, + {"id":"pkg/foo.go::Ghost","meta":{"summary":"no such node"}} + ]` + + res := callHandler(t, s.handleAnnotateNodes, map[string]any{"annotations": annotations}) + out := unmarshalResult(t, res) + t.Logf("[annotate#1] summary=%v", out) + + assert.Equal(t, 2.0, out["annotated"], "two real nodes annotated") + assert.Equal(t, 0.0, out["unchanged"]) + assert.Equal(t, 0.0, out["edges_added"]) + notFound, _ := out["not_found"].([]any) + require.Len(t, notFound, 1) + assert.Equal(t, "pkg/foo.go::Ghost", notFound[0]) + + // Round-trip: GetNode reflects the merged keys under the default + // "ext_" namespace. + bar := s.graph.GetNode("pkg/foo.go::Bar") + require.NotNil(t, bar) + assert.Equal(t, "parses the bar", bar.Meta["ext_summary"]) + assert.Equal(t, "ingest", bar.Meta["ext_domain"]) + assert.Equal(t, 0.7, bar.Meta["ext_complexity"]) + assert.Equal(t, []any{"parser", "io"}, bar.Meta["ext_tags"]) + // Structural fields untouched (MUST NOT). + assert.Equal(t, "Bar", bar.Name) + assert.Equal(t, graph.KindFunction, bar.Kind) + assert.Equal(t, 1, bar.StartLine) + assert.Equal(t, 9, bar.EndLine) + + // Idempotent re-run: same batch -> the two real nodes are unchanged. + res2 := callHandler(t, s.handleAnnotateNodes, map[string]any{"annotations": annotations}) + out2 := unmarshalResult(t, res2) + t.Logf("[annotate#2 idempotent] summary=%v", out2) + assert.Equal(t, 0.0, out2["annotated"], "re-run annotates nothing new") + assert.Equal(t, 2.0, out2["unchanged"], "both real nodes report unchanged") +} + +// TestHandleAnnotateNodes_Namespace proves the namespace argument drives +// the key prefix, and that a key already carrying the prefix is not +// double-prefixed. +func TestHandleAnnotateNodes_Namespace(t *testing.T) { + s := newAnnotateTestServer(t) + + res := callHandler(t, s.handleAnnotateNodes, map[string]any{ + "namespace": "ua", + "annotations": `[{"id":"pkg/foo.go::Bar","meta":{"summary":"x","ua_domain":"core"}}]`, + }) + out := unmarshalResult(t, res) + assert.Equal(t, 1.0, out["annotated"]) + + bar := s.graph.GetNode("pkg/foo.go::Bar") + require.NotNil(t, bar) + assert.Equal(t, "x", bar.Meta["ua_summary"], "bare key gets the namespace prefix") + assert.Equal(t, "core", bar.Meta["ua_domain"], "already-prefixed key is not double-prefixed") + _, doubled := bar.Meta["ua_ua_domain"] + assert.False(t, doubled, "must not double-prefix") +} + +// TestHandleAnnotateNodes_AddRelatedIdempotentEdge proves the optional +// add_related pairs create an idempotent semantically_related edge with +// the expected origin/confidence/similarity meta. +func TestHandleAnnotateNodes_AddRelatedIdempotentEdge(t *testing.T) { + s := newAnnotateTestServer(t) + + args := map[string]any{ + "annotations": `[{"id":"pkg/foo.go::Bar","meta":{"summary":"x"}}]`, + "add_related": `[["pkg/foo.go::Bar","pkg/foo.go::Baz",0.83]]`, + } + res := callHandler(t, s.handleAnnotateNodes, args) + out := unmarshalResult(t, res) + t.Logf("[related#1] summary=%v", out) + assert.Equal(t, 1.0, out["edges_added"]) + + // The edge exists with the expected shape. + outEdges := s.graph.GetOutEdges("pkg/foo.go::Bar") + var rel *graph.Edge + for _, e := range outEdges { + if e.Kind == graph.EdgeSemanticallyRelated && e.To == "pkg/foo.go::Baz" { + rel = e + break + } + } + require.NotNil(t, rel, "semantically_related edge should be present") + assert.Equal(t, "ext_annotated", rel.Origin) + assert.InDelta(t, 0.83, rel.Confidence, 1e-9) + assert.InDelta(t, 0.83, rel.Meta["similarity"].(float64), 1e-9) + + // Idempotent: re-adding the same pair does not duplicate the edge. + callHandler(t, s.handleAnnotateNodes, args) + count := 0 + for _, e := range s.graph.GetOutEdges("pkg/foo.go::Bar") { + if e.Kind == graph.EdgeSemanticallyRelated && e.To == "pkg/foo.go::Baz" { + count++ + } + } + assert.Equal(t, 1, count, "AddEdge dedup keeps a single semantically_related edge") +} + +// TestHandleAnnotateNodes_DefaultScore proves a pair without an explicit +// score defaults to 0.5. +func TestHandleAnnotateNodes_DefaultScore(t *testing.T) { + s := newAnnotateTestServer(t) + callHandler(t, s.handleAnnotateNodes, map[string]any{ + "annotations": `[{"id":"pkg/foo.go::Bar","meta":{"summary":"x"}}]`, + "add_related": `[["pkg/foo.go::Bar","pkg/foo.go::Baz"]]`, + }) + for _, e := range s.graph.GetOutEdges("pkg/foo.go::Bar") { + if e.Kind == graph.EdgeSemanticallyRelated { + assert.InDelta(t, 0.5, e.Confidence, 1e-9, "missing score defaults to 0.5") + } + } +} + +// TestHandleAnnotateNodes_BadInput proves malformed JSON and a missing +// required arg surface a clean error result (not a panic). +func TestHandleAnnotateNodes_BadInput(t *testing.T) { + s := newAnnotateTestServer(t) + + missing, err := s.handleAnnotateNodes(context.Background(), reqWith(map[string]any{})) + require.NoError(t, err) + assert.True(t, missing.IsError, "missing 'annotations' is an error result") + + bad, err := s.handleAnnotateNodes(context.Background(), reqWith(map[string]any{"annotations": "{not json"})) + require.NoError(t, err) + assert.True(t, bad.IsError, "malformed annotations JSON is an error result") +} + +// TestAnnotate_RegisteredOnNewServer guards against accidental removal +// of the registerAnnotateTools() call: the tool must be listed by a +// Server built through the real NewServer (registered + callable, and +// thus HTTP-exposed via the shared registry). +func TestAnnotate_RegisteredOnNewServer(t *testing.T) { + srv, _ := setupTestServer(t) + require.Contains(t, srv.mcpServer.ListTools(), "annotate_nodes", + "annotate_nodes must be registered on a NewServer-built server") +} + +// reqWith builds a CallToolRequest carrying the given arguments. +func reqWith(args map[string]any) mcp.CallToolRequest { + req := mcp.CallToolRequest{} + req.Params.Arguments = args + return req +}