From 2ec8b2c7fbf309ea408e21893d7ea5d40325a711 Mon Sep 17 00:00:00 2001
From: sgang <shreyansh.gang@nokia.com>
Date: Wed, 11 Mar 2026 10:13:11 +0530
Subject: [PATCH 01/10] serach a path with augmented path

---
 pkg/schema/container.go                     |  20 ++
 pkg/schema/expand.go                        |   6 +-
 pkg/schema/object.go                        |  78 +++++--
 pkg/schema/references.go                    |  15 ++
 pkg/schema/schema.go                        |  18 +-
 pkg/store/persiststore/persiststore.go      | 104 +++++----
 pkg/store/persiststore/persiststore_test.go | 246 ++++++++++++++++++++
 pkg/store/persiststore/yang_helpers.go      |  95 ++++++++
 8 files changed, 515 insertions(+), 67 deletions(-)
 create mode 100644 pkg/store/persiststore/persiststore_test.go
 create mode 100644 pkg/store/persiststore/yang_helpers.go

diff --git a/pkg/schema/container.go b/pkg/schema/container.go
index 6342f3e..7d9d933 100644
--- a/pkg/schema/container.go
+++ b/pkg/schema/container.go
@@ -108,6 +108,16 @@ func getMandatoryChildren(e *yang.Entry) []*sdcpb.MandatoryChild {
 			result = append(result, c)
 		}
 	}
+	// include mandatory children from augments
+	for _, v := range e.Augmented {
+		if v.Mandatory == yang.TSTrue {
+			c := &sdcpb.MandatoryChild{
+				Name:    v.Name,
+				IsState: isState(v),
+			}
+			result = append(result, c)
+		}
+	}
 	return result
 }
 
@@ -157,6 +167,16 @@ func getChoiceInfo(e *yang.Entry) *sdcpb.ChoiceInfo {
 
 		processChoice(de, ci)
 	}
+	// also consider choices introduced via augments
+	for _, de := range e.Augmented {
+		if !de.IsChoice() {
+			continue
+		}
+		if ci == nil {
+			ci = &sdcpb.ChoiceInfo{Choice: map[string]*sdcpb.ChoiceInfoChoice{}}
+		}
+		processChoice(de, ci)
+	}
 	return ci
 }
 
diff --git a/pkg/schema/expand.go b/pkg/schema/expand.go
index 7a9a09e..8527277 100644
--- a/pkg/schema/expand.go
+++ b/pkg/schema/expand.go
@@ -41,7 +41,7 @@ func (sc *Schema) ExpandPath(p *sdcpb.Path, dt sdcpb.DataType) ([]*sdcpb.Path, e
 	for _, k := range strings.Fields(e.Key) {
 		keys[k] = struct{}{}
 	}
-	for _, c := range e.Dir {
+	for _, c := range getChildren(e) {
 		// skip keys
 		if _, ok := keys[c.Name]; ok {
 			continue
@@ -109,7 +109,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 			kmap[k] = struct{}{}
 		}
 
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			if _, ok := kmap[c.Name]; ok {
 				continue
 			}
@@ -126,7 +126,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 	case e.IsContainer():
 		log.Debugf("got container: %s", e.Name)
 		containerPE := &sdcpb.PathElem{Name: e.Name, Key: make(map[string]string)}
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			log.Debugf("container parent adding child: %s", c.Name)
 			childrenPE := sc.getPathElems(c, dt)
 
diff --git a/pkg/schema/object.go b/pkg/schema/object.go
index df34983..4a5bf32 100644
--- a/pkg/schema/object.go
+++ b/pkg/schema/object.go
@@ -182,7 +182,7 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 		}
 		return e, nil
 	default:
-		if e.Dir == nil {
+		if e.Dir == nil && len(e.Augmented) == 0 {
 			return nil, errors.New("not found")
 		}
 		for _, ee := range getChildren(e) {
@@ -261,6 +261,22 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		return nil
 	}
 
+	// helper to find a direct child by name from Dir or Augmented
+	childByName := func(parent *yang.Entry, name string) *yang.Entry {
+		if parent == nil {
+			return nil
+		}
+		if ce, ok := parent.Dir[name]; ok {
+			return ce
+		}
+		for _, ae := range parent.Augmented {
+			if ae.Name == name {
+				return ae
+			}
+		}
+		return nil
+	}
+
 	switch {
 	case e.IsList():
 		if cpe.GetKey() == nil {
@@ -281,7 +297,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 			return nil
 		}
 		nxt := pe[count]
-		if ee, ok := e.Dir[nxt]; ok {
+		if ee := childByName(e, nxt); ee != nil {
 			return sc.buildPath(pe[count:], p, ee)
 		}
 		// find choices/cases
@@ -316,20 +332,20 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		return fmt.Errorf("case %s - unknown element %s", e.Name, pe[0])
 	case e.IsContainer():
 		// implicit case: child with same name which is a choice
-		if ee, ok := e.Dir[pe[0]]; ee != nil && ok {
+		if ee := childByName(e, pe[0]); ee != nil {
 			if ee.IsChoice() {
 				return sc.buildPath(pe[1:], p, ee)
 			}
 		}
 
 		p.Elem = append(p.Elem, cpe)
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := childByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		if lpe == 1 {
 			return nil
 		}
-		if ee, ok := e.Dir[pe[1]]; ok {
+		if ee := childByName(e, pe[1]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
 		// find choice/case
@@ -359,7 +375,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 func getChildren(e *yang.Entry) []*yang.Entry {
 	switch {
 	case e.IsChoice(), e.IsCase(), e.IsContainer(), e.IsList():
-		rs := make([]*yang.Entry, 0, len(e.Dir))
+		rs := make([]*yang.Entry, 0, len(e.Dir)+len(e.Augmented))
 		for _, ee := range e.Dir {
 			if ee.IsChoice() || ee.IsCase() {
 				rs = append(rs, getChildren(ee)...)
@@ -367,6 +383,14 @@ func getChildren(e *yang.Entry) []*yang.Entry {
 			}
 			rs = append(rs, ee)
 		}
+		// add augmented children as well
+		for _, ee := range e.Augmented {
+			if ee.IsChoice() || ee.IsCase() {
+				rs = append(rs, getChildren(ee)...)
+				continue
+			}
+			rs = append(rs, ee)
+		}
 		//sort.Slice(rs, sortFn(rs))
 		return rs
 		// case e.IsCase():
@@ -562,21 +586,41 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 	if len(pe) == 0 {
 		return e, nil
 	}
-	for _, ee := range e.Dir {
-		if !ee.IsChoice() {
-			continue
+	// pe is expected to contain at least the current element name at index 0 and
+	// the sought child element name at index 1.
+	//
+	// This is used from container/list resolution paths where the next element may
+	// live under a choice/case. Case nodes do not exist in the data tree, so we
+	// must search through cases to find the actual schema node.
+	if len(pe) < 2 {
+		return nil, fmt.Errorf("unknown element %s", pe[0])
+	}
+
+	// scan choice nodes in both direct and augmented children
+	choices := make([]*yang.Entry, 0)
+	for _, child := range e.Dir {
+		if child != nil && child.IsChoice() {
+			choices = append(choices, child)
 		}
-		if eee, ok := ee.Dir[pe[1]]; ok && !eee.IsCase() {
-			return eee, nil
+	}
+	for _, child := range e.Augmented {
+		if child != nil && child.IsChoice() {
+			choices = append(choices, child)
 		}
-		// assume there was a case obj,
-		// search one step deeper
-		for _, eee := range ee.Dir {
-			if !eee.IsCase() {
+	}
+
+	for _, choice := range choices {
+		// implicit case: choice directly contains the data node
+		if direct, ok := choice.Dir[pe[1]]; ok && direct != nil && !direct.IsCase() {
+			return direct, nil
+		}
+		// explicit cases: the data node is under a case
+		for _, cc := range choice.Dir {
+			if cc == nil || !cc.IsCase() {
 				continue
 			}
-			if eeee, ok := eee.Dir[pe[1]]; ok {
-				return eeee, nil
+			if target, ok := cc.Dir[pe[1]]; ok && target != nil {
+				return target, nil
 			}
 		}
 	}
diff --git a/pkg/schema/references.go b/pkg/schema/references.go
index a3f646a..3a1e6d7 100644
--- a/pkg/schema/references.go
+++ b/pkg/schema/references.go
@@ -70,6 +70,21 @@ func (sc *Schema) buildReferences(e *yang.Entry) error {
 			return err
 		}
 	}
+	// also recurse into augmented children
+	for _, ce := range e.Augmented {
+		if ce.IsCase() || ce.IsChoice() {
+			for _, cce := range ce.Dir {
+				err := sc.buildReferences(cce)
+				if err != nil {
+					return err
+				}
+			}
+		}
+		err := sc.buildReferences(ce)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
index 9ab9de0..dbeea39 100644
--- a/pkg/schema/schema.go
+++ b/pkg/schema/schema.go
@@ -141,14 +141,28 @@ func (s *Schema) Walk(e *yang.Entry, fn func(ec *yang.Entry) error) error {
 				return err
 			}
 		}
+		// also walk augments merged into this entry
+		for _, ee := range e.Augmented {
+			err = s.Walk(ee, fn)
+			if err != nil {
+				return err
+			}
+		}
 		return nil
 	}
 	err = fn(e)
 	if err != nil {
 		return err
 	}
-	for _, e := range e.Dir {
-		err = s.Walk(e, fn)
+	for _, ce := range e.Dir {
+		err = s.Walk(ce, fn)
+		if err != nil {
+			return err
+		}
+	}
+	// include augmented children at every level
+	for _, ce := range e.Augmented {
+		err = s.Walk(ce, fn)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index 2bad4c7..2d83fec 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -733,6 +733,8 @@ func getModules(txn *badger.Txn, sc store.SchemaKey) ([]string, error) {
 
 func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest, sck store.SchemaKey) (*sdcpb.GetSchemaResponse, error) {
 	pes := utils.ToStrings(req.GetPath(), false, true)
+	log.Debugf("[persiststore][getSchema] raw path elems=%v", pes)
+
 	cKey := cacheKey{
 		SchemaKey: sck,
 		Path:      strings.Join(pes, "/"),
@@ -750,92 +752,104 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 			return rsp, nil
 		}
 	}
+
 	var err error
 	sce := new(sdcpb.SchemaElem)
+	var modules []string
 
-	// key all i.e "root"
-	if lpes := len(pes); lpes == 0 || (lpes == 1 && pes[0] == "") {
-		err = s.db.View(func(txn *badger.Txn) error {
+	origin := req.GetPath().GetOrigin()
+
+	// Root schema
+	if len(pes) == 0 || (len(pes) == 1 && pes[0] == "") {
+		var sce sdcpb.SchemaElem
+		err := s.db.View(func(txn *badger.Txn) error {
 			k := buildEntryKey(sck, []string{schema.RootName})
 			item, err := txn.Get(k)
 			if err != nil {
 				return err
 			}
-			if item == nil {
-				return ErrKeyNotFound
-			}
 			v, err := item.ValueCopy(nil)
 			if err != nil {
 				return err
 			}
-			err = proto.Unmarshal(v, sce)
-			if err != nil {
-				return err
-			}
-			return nil
+			return proto.Unmarshal(v, &sce)
 		})
 		if err != nil {
 			return nil, err
 		}
-		return &sdcpb.GetSchemaResponse{Schema: sce}, nil
+		return &sdcpb.GetSchemaResponse{Schema: &sce}, nil
 	}
-	moduleName := ""
-	if index := strings.Index(pes[0], ":"); index > 0 {
-		moduleName = pes[0][:index]
-		pes[0] = pes[0][index+1:]
+
+	// Parse per-element prefixes and build unprefixed names
+	pp := parsePathElems(pes)
+	names := make([]string, 0, len(pp))
+	for _, pe := range pp {
+		names = append(names, pe.name)
 	}
-	var modules []string
-	// path has module prefix
-	if moduleName != "" {
-		modules = []string{moduleName}
-	} else {
-		// path does not have module prefix
-		modules, err = s.getModules(sck)
-		if err != nil {
-			return nil, err
+
+	// Apply gNMI origin as module hint on first element (if present)
+	if origin != "" && len(pp) > 0 && pp[0].module == "" {
+		pp[0].module = origin
+	}
+
+	log.Debugf("[persiststore][getSchema] parsed path elems=%+v", pp)
+
+	// Validate first element's module prefix if present
+	// Note: non-first element prefixes can refer to augmented modules not in root list
+	allModules, err := s.getModules(sck)
+	if err != nil {
+		return nil, err
+	}
+	if len(pp) > 0 && pp[0].module != "" {
+		modSet := make(map[string]struct{}, len(allModules))
+		for _, m := range allModules {
+			modSet[m] = struct{}{}
+		}
+		if _, ok := modSet[pp[0].module]; !ok {
+			return nil, status.Errorf(codes.InvalidArgument, "unknown module prefix %q", pp[0].module)
+		}
+		if pp[0].name == "" {
+			return nil, status.Errorf(codes.InvalidArgument, "empty identifier after prefix %q", pp[0].module)
 		}
+	}
+
+	// Decide candidate modules: use first element's module if present, else try all
+	if len(pp) > 0 && pp[0].module != "" {
+		modules = []string{pp[0].module}
+	} else {
+		// Prefer modules in a stable, deprioritized order
+		modules = append(modules, allModules...)
 		sort.Slice(modules, func(i, j int) bool {
 			return utils.SortModulesAB(modules[i], modules[j], config.DeprioritizedModules)
 		})
 	}
 
-	npe := make([]string, 1+len(pes))
-	copy(npe[1:], pes)
 	err = s.db.View(func(txn *badger.Txn) error {
 		for _, module := range modules {
-			var k []byte
-			if npe[1] == module { // query module name
-				k = buildEntryKey(sck, npe[1:])
-			} else {
-				npe[0] = module
-				k = buildEntryKey(sck, npe)
-			}
+			keyPath := make([]string, 0, 1+len(names))
+			keyPath = append(keyPath, module)
+			keyPath = append(keyPath, names...)
+			k := buildEntryKey(sck, keyPath)
+
 			item, err := txn.Get(k)
-			if err != nil {
-				continue
-			}
-			if item == nil {
+			if err != nil || item == nil {
 				continue
 			}
 			v, err := item.ValueCopy(nil)
 			if err != nil {
 				return err
 			}
-			err = proto.Unmarshal(v, sce)
-			if err != nil {
+			if err := proto.Unmarshal(v, sce); err != nil {
 				return err
 			}
 			return nil
 		}
-		return fmt.Errorf("%s: %w", req.GetPath(), ErrKeyNotFound)
+		return fmt.Errorf("schema path not found: %s", req.GetPath())
 	})
 	if err != nil {
 		return nil, err
 	}
-	rsp := &sdcpb.GetSchemaResponse{Schema: sce}
-	if s.cache != nil {
-		s.cache.Set(cKey, rsp, ttlcache.DefaultTTL)
-	}
+
 	return &sdcpb.GetSchemaResponse{Schema: sce}, nil
 }
 
diff --git a/pkg/store/persiststore/persiststore_test.go b/pkg/store/persiststore/persiststore_test.go
new file mode 100644
index 0000000..cdc217f
--- /dev/null
+++ b/pkg/store/persiststore/persiststore_test.go
@@ -0,0 +1,246 @@
+// Copyright 2024 Nokia
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package persiststore
+
+import (
+	"context"
+	"testing"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/sdcio/schema-server/pkg/schema"
+	"github.com/sdcio/schema-server/pkg/store"
+	sdcpb "github.com/sdcio/sdc-protos/sdcpb"
+	"google.golang.org/protobuf/proto"
+)
+
+//
+// ---------- helpers ----------
+//
+
+func newTestStore(t *testing.T) *persistStore {
+	t.Helper()
+
+	dir := t.TempDir()
+	db, err := badger.Open(badger.DefaultOptions(dir).WithLogger(nil))
+	if err != nil {
+		t.Fatalf("failed to open badger: %v", err)
+	}
+
+	t.Cleanup(func() { _ = db.Close() })
+
+	return &persistStore{db: db}
+}
+
+func testSchemaKey() store.SchemaKey {
+	return store.SchemaKey{
+		Name:    "M",
+		Vendor:  "V",
+		Version: "1",
+	}
+}
+
+func insertSchemaMeta(t *testing.T, ps *persistStore, sk store.SchemaKey) {
+	t.Helper()
+
+	key := buildSchemaKey(sk)
+	err := ps.db.Update(func(txn *badger.Txn) error {
+		return txn.Set(key, []byte(`{}`))
+	})
+	if err != nil {
+		t.Fatalf("failed inserting schema meta: %v", err)
+	}
+}
+
+func insertRootEntry(t *testing.T, ps *persistStore, sk store.SchemaKey, modules []string) {
+	t.Helper()
+
+	// Create a root container listing available modules as children
+	root := &sdcpb.SchemaElem{
+		Schema: &sdcpb.SchemaElem_Container{Container: &sdcpb.ContainerSchema{
+			Name:     schema.RootName,
+			Children: modules,
+		}},
+	}
+	b, err := proto.Marshal(root)
+	if err != nil {
+		t.Fatalf("marshal root: %v", err)
+	}
+	key := buildEntryKey(sk, []string{schema.RootName})
+	if err := ps.db.Update(func(txn *badger.Txn) error { return txn.Set(key, b) }); err != nil {
+		t.Fatalf("insert root failed: %v", err)
+	}
+}
+
+func insertEntry(t *testing.T, ps *persistStore, sk store.SchemaKey, keyPath []string, se *sdcpb.SchemaElem) {
+	t.Helper()
+	b, err := proto.Marshal(se)
+	if err != nil {
+		t.Fatalf("marshal entry: %v", err)
+	}
+	key := buildEntryKey(sk, keyPath)
+	if err := ps.db.Update(func(txn *badger.Txn) error { return txn.Set(key, b) }); err != nil {
+		t.Fatalf("insert entry failed: %v", err)
+	}
+}
+
+//
+// ---------- helper function tests ----------
+//
+
+func TestSchemaKeyString(t *testing.T) {
+	sk := store.SchemaKey{Name: "n", Vendor: "v", Version: "1"}
+	if got := schemaKeyString(sk); got != "n@v@1" {
+		t.Fatalf("unexpected schemaKeyString: %q", got)
+	}
+}
+
+func TestStripPrefix(t *testing.T) {
+	cases := map[string]string{
+		"a":       "a",
+		"m:a":     "a",
+		"foo:bar": "bar",
+	}
+
+	for in, exp := range cases {
+		if got := stripPrefix(in); got != exp {
+			t.Fatalf("stripPrefix(%q)=%q, want %q", in, got, exp)
+		}
+	}
+}
+
+func TestHasPrefix(t *testing.T) {
+	if !hasPrefix("m:a") {
+		t.Fatalf("expected prefix")
+	}
+	if hasPrefix("a") {
+		t.Fatalf("unexpected prefix")
+	}
+}
+
+//
+// ---------- HasSchema tests ----------
+//
+
+func TestHasSchema(t *testing.T) {
+	ps := newTestStore(t)
+	sk := testSchemaKey()
+
+	if ps.HasSchema(sk) {
+		t.Fatalf("schema should not exist")
+	}
+
+	insertSchemaMeta(t, ps, sk)
+
+	if !ps.HasSchema(sk) {
+		t.Fatalf("schema should exist")
+	}
+}
+
+//
+// ---------- GetSchema tests (negative + strict) ----------
+//
+
+func TestGetSchema_UnknownSchema(t *testing.T) {
+	ps := newTestStore(t)
+
+	_, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{
+			Name:    "M",
+			Vendor:  "V",
+			Version: "1",
+		},
+	})
+	if err == nil {
+		t.Fatalf("expected error for unknown schema")
+	}
+}
+
+func TestGetSchema_StrictPrefixRejected(t *testing.T) {
+	// Replace strict-prefix behavior with validation of unknown module hints
+	ps := newTestStore(t)
+	sk := testSchemaKey()
+	insertSchemaMeta(t, ps, sk)
+	// Insert root with one known module
+	insertRootEntry(t, ps, sk, []string{"known"})
+
+	_, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Elem: []*sdcpb.PathElem{{Name: "unknown:foo"}}},
+	})
+	if err == nil {
+		t.Fatalf("expected error for unknown module prefix")
+	}
+}
+
+func TestGetSchema_RootSchemaMissingEntry(t *testing.T) {
+	ps := newTestStore(t)
+	sk := testSchemaKey()
+	insertSchemaMeta(t, ps, sk)
+
+	_, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{
+			Name:    sk.Name,
+			Vendor:  sk.Vendor,
+			Version: sk.Version,
+		},
+	})
+	if err == nil {
+		t.Fatalf("expected error due to missing root entry")
+	}
+}
+
+func TestGetSchema_ModuleLessPathResolves(t *testing.T) {
+	ps := newTestStore(t)
+	sk := testSchemaKey()
+	insertSchemaMeta(t, ps, sk)
+	// Set up root with module list and a concrete entry under that module
+	insertRootEntry(t, ps, sk, []string{"ietf-nss"})
+	// Create container entry for ietf-nss:network-instances
+	se := &sdcpb.SchemaElem{Schema: &sdcpb.SchemaElem_Container{Container: &sdcpb.ContainerSchema{Name: "network-instances"}}}
+	insertEntry(t, ps, sk, []string{"ietf-nss", "network-instances"}, se)
+
+	rsp, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Elem: []*sdcpb.PathElem{{Name: "network-instances"}}},
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	got := rsp.GetSchema().GetContainer().GetName()
+	if got != "network-instances" {
+		t.Fatalf("unexpected schema name: %q", got)
+	}
+}
+
+func TestGetSchema_ModulePrefixedPathResolves(t *testing.T) {
+	ps := newTestStore(t)
+	sk := testSchemaKey()
+	insertSchemaMeta(t, ps, sk)
+	insertRootEntry(t, ps, sk, []string{"ietf-nss"})
+	se := &sdcpb.SchemaElem{Schema: &sdcpb.SchemaElem_Container{Container: &sdcpb.ContainerSchema{Name: "network-instances"}}}
+	insertEntry(t, ps, sk, []string{"ietf-nss", "network-instances"}, se)
+
+	rsp, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Elem: []*sdcpb.PathElem{{Name: "ietf-nss:network-instances"}}},
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	got := rsp.GetSchema().GetContainer().GetName()
+	if got != "network-instances" {
+		t.Fatalf("unexpected schema name: %q", got)
+	}
+}
diff --git a/pkg/store/persiststore/yang_helpers.go b/pkg/store/persiststore/yang_helpers.go
new file mode 100644
index 0000000..d54f120
--- /dev/null
+++ b/pkg/store/persiststore/yang_helpers.go
@@ -0,0 +1,95 @@
+package persiststore
+
+import (
+	"strings"
+
+	sdcpb "github.com/sdcio/sdc-protos/sdcpb"
+)
+
+// pathElem represents a parsed path element with an optional module hint.
+//
+// The module is derived from a prefix in the form "<module>:<name>" when present.
+// The name field always contains the unprefixed element name.
+type pathElem struct {
+	name   string // unprefixed name
+	module string // optional module hint
+}
+
+// parsePathElems parses gNMI-like path elements that may optionally carry a module
+// prefix in the form "<module>:<name>".
+//
+// It returns a slice of pathElem values where:
+// - name is always the unprefixed element name
+// - module is set only when a prefix was present
+//
+// Note: only the first ':' is treated as the prefix separator.
+
+func parsePathElems(pes []string) []pathElem {
+	out := make([]pathElem, 0, len(pes))
+	for _, pe := range pes {
+		if i := strings.IndexByte(pe, ':'); i > 0 {
+			out = append(out, pathElem{
+				module: pe[:i],
+				name:   pe[i+1:],
+			})
+		} else {
+			out = append(out, pathElem{
+				name: pe,
+			})
+		}
+	}
+	return out
+}
+
+// uniqueStrings returns the input slice with duplicates removed while preserving
+// the original order of first occurrence.
+
+func uniqueStrings(in []string) []string {
+	m := make(map[string]struct{})
+	out := make([]string, 0, len(in))
+	for _, s := range in {
+		if _, ok := m[s]; !ok {
+			m[s] = struct{}{}
+			out = append(out, s)
+		}
+	}
+	return out
+}
+
+// schemaElemModuleName extracts the module name from a SchemaElem, regardless of
+// whether the element is a container, leaf, or leaf-list.
+//
+// If the schema element is nil or of an unknown/unsupported oneof type, an empty
+// string is returned.
+
+func schemaElemModuleName(sce *sdcpb.SchemaElem) string {
+	switch s := sce.Schema.(type) {
+	case *sdcpb.SchemaElem_Container:
+		return s.Container.ModuleName
+	case *sdcpb.SchemaElem_Leaflist:
+		return s.Leaflist.ModuleName
+	case *sdcpb.SchemaElem_Field:
+		return s.Field.ModuleName
+	default:
+		return ""
+	}
+}
+
+// hasPrefix reports whether the provided path element contains a module prefix.
+// A prefix is identified by the presence of ':' anywhere in the string.
+
+func hasPrefix(pe string) bool {
+	return strings.Contains(pe, ":")
+}
+
+// stripPrefix removes the module prefix from a path element of the form
+// "<module>:<name>".
+//
+// If no ':' is present, the input is returned unchanged.
+
+func stripPrefix(pe string) string {
+	if i := strings.IndexByte(pe, ':'); i != -1 {
+		return pe[i+1:]
+	}
+	return pe
+}

From f0b7cc4a6c45e93f24e13cb94ed0dd373e248b1a Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 10:43:09 +0200
Subject: [PATCH 02/10] bring back the cache population

---
 pkg/store/persiststore/persiststore.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index 2d83fec..efdc0f6 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -849,7 +849,10 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 	if err != nil {
 		return nil, err
 	}
-
+	rsp := &sdcpb.GetSchemaResponse{Schema: sce}
+	if s.cache != nil {
+		s.cache.Set(cKey, rsp, ttlcache.DefaultTTL)
+	}
 	return &sdcpb.GetSchemaResponse{Schema: sce}, nil
 }
 

From 06ad9937d96116ac188c9a3d5490e9a1153912c8 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 10:50:22 +0200
Subject: [PATCH 03/10] also cache root node

---
 pkg/store/persiststore/persiststore.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index efdc0f6..d62ad78 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -761,7 +761,7 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 
 	// Root schema
 	if len(pes) == 0 || (len(pes) == 1 && pes[0] == "") {
-		var sce sdcpb.SchemaElem
+		sce = new(sdcpb.SchemaElem)
 		err := s.db.View(func(txn *badger.Txn) error {
 			k := buildEntryKey(sck, []string{schema.RootName})
 			item, err := txn.Get(k)
@@ -772,12 +772,16 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 			if err != nil {
 				return err
 			}
-			return proto.Unmarshal(v, &sce)
+			return proto.Unmarshal(v, sce)
 		})
 		if err != nil {
 			return nil, err
 		}
-		return &sdcpb.GetSchemaResponse{Schema: &sce}, nil
+		rsp := &sdcpb.GetSchemaResponse{Schema: sce}
+		if s.cache != nil {
+			s.cache.Set(cKey, rsp, ttlcache.DefaultTTL)
+		}
+		return rsp, nil
 	}
 
 	// Parse per-element prefixes and build unprefixed names
@@ -853,7 +857,7 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 	if s.cache != nil {
 		s.cache.Set(cKey, rsp, ttlcache.DefaultTTL)
 	}
-	return &sdcpb.GetSchemaResponse{Schema: sce}, nil
+	return rsp, nil
 }
 
 func removeDescription(rsp *sdcpb.GetSchemaResponse) *sdcpb.GetSchemaResponse {

From 5ae27f7ad7f6df7bf54acd61210040d3ea571b45 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 11:02:23 +0200
Subject: [PATCH 04/10] fix multi-allocation

---
 pkg/store/persiststore/persiststore.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index d62ad78..e408362 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -761,7 +761,6 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 
 	// Root schema
 	if len(pes) == 0 || (len(pes) == 1 && pes[0] == "") {
-		sce = new(sdcpb.SchemaElem)
 		err := s.db.View(func(txn *badger.Txn) error {
 			k := buildEntryKey(sck, []string{schema.RootName})
 			item, err := txn.Get(k)

From 4fb26bae6bff65eb61863a1fa1bdf07b899ee20b Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 11:14:45 +0200
Subject: [PATCH 05/10] =?UTF-8?q?`persiststore.addSchemaElem`=20recurses?=
 =?UTF-8?q?=20**`e.Dir`=20only**,=20not=20**`e.Augmented`**=20=E2=86=92=20?=
 =?UTF-8?q?augment-only=20nodes=20may=20never=20hit=20Badger=20while=20`me?=
 =?UTF-8?q?mstore`=20can=20resolve=20them.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

addSchemaElem now walks Augmented (persiststore.go)
Choice/case: after recursing e.Dir, it also recurses e.Augmented (same as Schema.Walk).
Normal nodes: after persisting children from e.Dir, it does the same for e.Augmented, so augment-only nodes get Badger keys and protos like the rest of the tree.
---
 pkg/store/persiststore/persiststore.go        | 14 ++++++-
 pkg/store/persiststore/persiststore_test.go   | 38 +++++++++++++++++++
 .../testdata/augment/aug-base.yang            | 15 ++++++++
 .../testdata/augment/aug-extra.yang           | 18 +++++++++
 4 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 pkg/store/persiststore/testdata/augment/aug-base.yang
 create mode 100644 pkg/store/persiststore/testdata/augment/aug-extra.yang

diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index e408362..f9f497d 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -577,6 +577,12 @@ func (s *persistStore) addSchemaElem(wb *badger.WriteBatch, sc *schema.Schema, e
 				return err
 			}
 		}
+		for _, ee := range e.Augmented {
+			err := s.addSchemaElem(wb, sc, ee)
+			if err != nil {
+				return err
+			}
+		}
 		return nil
 	}
 	// build entry key
@@ -595,13 +601,19 @@ func (s *persistStore) addSchemaElem(wb *badger.WriteBatch, sc *schema.Schema, e
 	if err != nil {
 		return err
 	}
-	// do the same for entry children
+	// do the same for entry children (native definitions and augment-only nodes)
 	for _, ee := range e.Dir {
 		err = s.addSchemaElem(wb, sc, ee)
 		if err != nil {
 			return err
 		}
 	}
+	for _, ee := range e.Augmented {
+		err = s.addSchemaElem(wb, sc, ee)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
diff --git a/pkg/store/persiststore/persiststore_test.go b/pkg/store/persiststore/persiststore_test.go
index cdc217f..6bc6ab5 100644
--- a/pkg/store/persiststore/persiststore_test.go
+++ b/pkg/store/persiststore/persiststore_test.go
@@ -16,9 +16,11 @@ package persiststore
 
 import (
 	"context"
+	"path/filepath"
 	"testing"
 
 	"github.com/dgraph-io/badger/v4"
+	"github.com/sdcio/schema-server/pkg/config"
 	"github.com/sdcio/schema-server/pkg/schema"
 	"github.com/sdcio/schema-server/pkg/store"
 	sdcpb "github.com/sdcio/sdc-protos/sdcpb"
@@ -224,6 +226,42 @@ func TestGetSchema_ModuleLessPathResolves(t *testing.T) {
 	}
 }
 
+func TestAddSchema_PersistsAugmentOnlyChildren(t *testing.T) {
+	td := filepath.Join("testdata", "augment")
+	base := filepath.Join(td, "aug-base.yang")
+	extra := filepath.Join(td, "aug-extra.yang")
+
+	sc, err := schema.NewSchema(&config.SchemaConfig{
+		Name:    "augment-persist",
+		Vendor:  "test",
+		Version: "0",
+		Files:   []string{base, extra},
+	})
+	if err != nil {
+		t.Fatalf("NewSchema: %v", err)
+	}
+
+	ps := newTestStore(t)
+	if err := ps.AddSchema(sc); err != nil {
+		t.Fatalf("AddSchema: %v", err)
+	}
+
+	// augment-only-leaf exists only under target via aug-extra; must be readable from Badger after Reset.
+	rsp, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sc.Name(), Vendor: sc.Vendor(), Version: sc.Version()},
+		Path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+			{Name: "target"},
+			{Name: "augment-only-leaf"},
+		}},
+	})
+	if err != nil {
+		t.Fatalf("GetSchema augment-only path: %v", err)
+	}
+	if rsp.GetSchema().GetField().GetName() != "augment-only-leaf" {
+		t.Fatalf("expected field augment-only-leaf, got %#v", rsp.GetSchema())
+	}
+}
+
 func TestGetSchema_ModulePrefixedPathResolves(t *testing.T) {
 	ps := newTestStore(t)
 	sk := testSchemaKey()
diff --git a/pkg/store/persiststore/testdata/augment/aug-base.yang b/pkg/store/persiststore/testdata/augment/aug-base.yang
new file mode 100644
index 0000000..26390cd
--- /dev/null
+++ b/pkg/store/persiststore/testdata/augment/aug-base.yang
@@ -0,0 +1,15 @@
+module aug-base {
+  yang-version 1.1;
+  namespace "urn:augment:test:base";
+  prefix ab;
+
+  organization "schema-server test";
+  description "Base module for augment persistence test";
+
+  container target {
+    description "Augment target";
+    leaf native-leaf {
+      type string;
+    }
+  }
+}
diff --git a/pkg/store/persiststore/testdata/augment/aug-extra.yang b/pkg/store/persiststore/testdata/augment/aug-extra.yang
new file mode 100644
index 0000000..8a599e6
--- /dev/null
+++ b/pkg/store/persiststore/testdata/augment/aug-extra.yang
@@ -0,0 +1,18 @@
+module aug-extra {
+  yang-version 1.1;
+  namespace "urn:augment:test:extra";
+  prefix ax;
+
+  import aug-base {
+    prefix ab;
+  }
+
+  organization "schema-server test";
+  description "Adds augment-only leaf under aug-base:target";
+
+  augment "/ab:target" {
+    leaf augment-only-leaf {
+      type string;
+    }
+  }
+}

From 810a2ceb92c72fbd6a73c3898f73d56c66e14618 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 12:44:26 +0200
Subject: [PATCH 06/10] fix(schema,persiststore): persist and resolve
 augment-only schema nodes - Recurse yang.Entry.Augmented in
 persiststore.addSchemaElem (including   choice/case) so Badger seeding
 matches getChildren/Walk; add augment   regression test and testdata. - Add
 entryChildByName and use it in buildPath for list/container and  
 augment-aware choice/case branches; extend findChoiceCase for   implicit
 children and cases under choice.Augmented. - Add bpcc-base/bpcc-aug YANG and
 TestSchema_BuildPath_AugmentedUnderCase   plus native leaf-a control path.

---
 pkg/schema/object.go                          | 90 ++++++++++++-------
 pkg/schema/object_test.go                     | 41 +++++++++
 .../testdata/buildpath-augment/bpcc-aug.yang  | 18 ++++
 .../testdata/buildpath-augment/bpcc-base.yang | 23 +++++
 4 files changed, 139 insertions(+), 33 deletions(-)
 create mode 100644 pkg/schema/testdata/buildpath-augment/bpcc-aug.yang
 create mode 100644 pkg/schema/testdata/buildpath-augment/bpcc-base.yang

diff --git a/pkg/schema/object.go b/pkg/schema/object.go
index 4a5bf32..d5cb05f 100644
--- a/pkg/schema/object.go
+++ b/pkg/schema/object.go
@@ -201,6 +201,22 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 	}
 }
 
+// entryChildByName returns a direct child of parent by name, checking Dir then Augmented.
+func entryChildByName(parent *yang.Entry, name string) *yang.Entry {
+	if parent == nil {
+		return nil
+	}
+	if ce, ok := parent.Dir[name]; ok {
+		return ce
+	}
+	for _, ae := range parent.Augmented {
+		if ae != nil && ae.Name == name {
+			return ae
+		}
+	}
+	return nil
+}
+
 func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 	if len(pe) == 0 {
 		return nil
@@ -261,22 +277,6 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		return nil
 	}
 
-	// helper to find a direct child by name from Dir or Augmented
-	childByName := func(parent *yang.Entry, name string) *yang.Entry {
-		if parent == nil {
-			return nil
-		}
-		if ce, ok := parent.Dir[name]; ok {
-			return ce
-		}
-		for _, ae := range parent.Augmented {
-			if ae.Name == name {
-				return ae
-			}
-		}
-		return nil
-	}
-
 	switch {
 	case e.IsList():
 		if cpe.GetKey() == nil {
@@ -297,7 +297,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 			return nil
 		}
 		nxt := pe[count]
-		if ee := childByName(e, nxt); ee != nil {
+		if ee := entryChildByName(e, nxt); ee != nil {
 			return sc.buildPath(pe[count:], p, ee)
 		}
 		// find choices/cases
@@ -307,45 +307,61 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		}
 		return sc.buildPath(pe[count:], p, ee)
 	case e.IsChoice():
-		p.Elem = append(p.Elem, cpe)
-		for _, entry := range e.Dir {
-			if entry.IsCase() {
-				if ee, ok := entry.Dir[pe[0]]; ok {
-					return sc.buildPath(pe[1:], p, ee)
+		{
+			tryChoiceChild := func(choice *yang.Entry, entry *yang.Entry) (handled bool, err error) {
+				if entry == nil {
+					return false, nil
+				}
+				if entry.IsCase() {
+					if ee := entryChildByName(entry, pe[0]); ee != nil {
+						return true, sc.buildPath(pe[1:], p, ee)
+					}
+				} else {
+					if ee := entryChildByName(choice, pe[0]); ee != nil {
+						return true, sc.buildPath(pe[1:], p, ee)
+					}
 				}
-			} else {
-				if ee, ok := e.Dir[pe[0]]; ok {
-					return sc.buildPath(pe[1:], p, ee)
+				return false, nil
+			}
+			p.Elem = append(p.Elem, cpe)
+			for _, entry := range e.Dir {
+				if handled, err := tryChoiceChild(e, entry); handled {
+					return err
+				}
+			}
+			for _, entry := range e.Augmented {
+				if handled, err := tryChoiceChild(e, entry); handled {
+					return err
 				}
 			}
+			return fmt.Errorf("choice %s - unknown element %s", e.Name, pe[0])
 		}
-		return fmt.Errorf("choice %s - unknown element %s", e.Name, pe[0])
 	case e.IsCase():
 		// RFC7950 7.9.2: A case node does not exist in the data tree.
 		// p.Elem = append(p.Elem, cpe)
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
-		if ee, ok := e.Dir[e.Name]; ok {
+		if ee := entryChildByName(e, e.Name); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		return fmt.Errorf("case %s - unknown element %s", e.Name, pe[0])
 	case e.IsContainer():
 		// implicit case: child with same name which is a choice
-		if ee := childByName(e, pe[0]); ee != nil {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			if ee.IsChoice() {
 				return sc.buildPath(pe[1:], p, ee)
 			}
 		}
 
 		p.Elem = append(p.Elem, cpe)
-		if ee := childByName(e, pe[0]); ee != nil {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		if lpe == 1 {
 			return nil
 		}
-		if ee := childByName(e, pe[1]); ee != nil {
+		if ee := entryChildByName(e, pe[1]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
 		// find choice/case
@@ -611,7 +627,7 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 
 	for _, choice := range choices {
 		// implicit case: choice directly contains the data node
-		if direct, ok := choice.Dir[pe[1]]; ok && direct != nil && !direct.IsCase() {
+		if direct := entryChildByName(choice, pe[1]); direct != nil && !direct.IsCase() {
 			return direct, nil
 		}
 		// explicit cases: the data node is under a case
@@ -619,7 +635,15 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 			if cc == nil || !cc.IsCase() {
 				continue
 			}
-			if target, ok := cc.Dir[pe[1]]; ok && target != nil {
+			if target := entryChildByName(cc, pe[1]); target != nil {
+				return target, nil
+			}
+		}
+		for _, cc := range choice.Augmented {
+			if cc == nil || !cc.IsCase() {
+				continue
+			}
+			if target := entryChildByName(cc, pe[1]); target != nil {
 				return target, nil
 			}
 		}
diff --git a/pkg/schema/object_test.go b/pkg/schema/object_test.go
index f10ca05..df02ddf 100644
--- a/pkg/schema/object_test.go
+++ b/pkg/schema/object_test.go
@@ -15,6 +15,7 @@
 package schema
 
 import (
+	"path/filepath"
 	"reflect"
 	"sort"
 	"testing"
@@ -476,6 +477,46 @@ func TestSchema_BuildPath(t *testing.T) {
 	}
 }
 
+func TestSchema_BuildPath_AugmentedUnderCase(t *testing.T) {
+	td := filepath.Join("testdata", "buildpath-augment")
+	sc, err := NewSchema(&config.SchemaConfig{
+		Name:    "bpcc",
+		Vendor:  "test",
+		Version: "0",
+		Files: []string{
+			filepath.Join(td, "bpcc-base.yang"),
+			filepath.Join(td, "bpcc-aug.yang"),
+		},
+	})
+	if err != nil {
+		t.Fatalf("NewSchema: %v", err)
+	}
+	p := &sdcpb.Path{}
+	err = sc.BuildPath([]string{"bpcc-base:top", "augment-leaf"}, p)
+	if err != nil {
+		t.Fatalf("BuildPath: %v", err)
+	}
+	want := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "bpcc-base:top"},
+		{Name: "augment-leaf"},
+	}}
+	if !comparePaths(p, want) {
+		t.Fatalf("got %v want %v", p, want)
+	}
+
+	p2 := &sdcpb.Path{}
+	if err := sc.BuildPath([]string{"bpcc-base:top", "leaf-a"}, p2); err != nil {
+		t.Fatalf("BuildPath native leaf under case: %v", err)
+	}
+	want2 := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "bpcc-base:top"},
+		{Name: "leaf-a"},
+	}}
+	if !comparePaths(p2, want2) {
+		t.Fatalf("native choice path got %v want %v", p2, want2)
+	}
+}
+
 func comparePathElem(pe1, pe2 *sdcpb.PathElem) bool {
 	if pe1 == nil {
 		return pe2 == nil
diff --git a/pkg/schema/testdata/buildpath-augment/bpcc-aug.yang b/pkg/schema/testdata/buildpath-augment/bpcc-aug.yang
new file mode 100644
index 0000000..de0b0da
--- /dev/null
+++ b/pkg/schema/testdata/buildpath-augment/bpcc-aug.yang
@@ -0,0 +1,18 @@
+module bpcc-aug {
+  yang-version 1.1;
+  namespace "urn:bpcc:aug";
+  prefix bcx;
+
+  import bpcc-base {
+    prefix bcb;
+  }
+
+  organization "schema-server test";
+  description "Augment-only leaf under case-a";
+
+  augment "/bcb:top/bcb:variant/bcb:case-a" {
+    leaf augment-leaf {
+      type string;
+    }
+  }
+}
diff --git a/pkg/schema/testdata/buildpath-augment/bpcc-base.yang b/pkg/schema/testdata/buildpath-augment/bpcc-base.yang
new file mode 100644
index 0000000..7a7be86
--- /dev/null
+++ b/pkg/schema/testdata/buildpath-augment/bpcc-base.yang
@@ -0,0 +1,23 @@
+module bpcc-base {
+  yang-version 1.1;
+  namespace "urn:bpcc:base";
+  prefix bcb;
+
+  organization "schema-server test";
+  description "Base module: choice/case for BuildPath augment tests";
+
+  container top {
+    choice variant {
+      case case-a {
+        leaf leaf-a {
+          type string;
+        }
+      }
+      case case-b {
+        leaf leaf-b {
+          type string;
+        }
+      }
+    }
+  }
+}

From 7290b8b7dbe5aca047b51c27f12db6101c222491 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 13:07:11 +0200
Subject: [PATCH 07/10] fix(persiststore): include Path.origin in GetSchema
 cache key

ToStrings(..., prefix=false) omits origin, but origin drives the module
hint for unprefixed paths, so identical path strings could cache the
wrong SchemaElem. Add Origin to cacheKey and a regression test with the
TTL cache enabled.
---
 pkg/store/persiststore/persiststore.go      |  7 ++-
 pkg/store/persiststore/persiststore_test.go | 64 +++++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index f9f497d..0b7c0b4 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -54,7 +54,8 @@ var (
 
 type cacheKey struct {
 	store.SchemaKey
-	Path string
+	Path   string
+	Origin string
 }
 
 type persistStore struct {
@@ -747,9 +748,11 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 	pes := utils.ToStrings(req.GetPath(), false, true)
 	log.Debugf("[persiststore][getSchema] raw path elems=%v", pes)
 
+	origin := req.GetPath().GetOrigin()
 	cKey := cacheKey{
 		SchemaKey: sck,
 		Path:      strings.Join(pes, "/"),
+		Origin:    origin,
 	}
 	if s.cache != nil {
 		if item := s.cache.Get(cKey, ttlcache.WithDisableTouchOnHit[cacheKey, *sdcpb.GetSchemaResponse]()); item != nil {
@@ -769,8 +772,6 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 	sce := new(sdcpb.SchemaElem)
 	var modules []string
 
-	origin := req.GetPath().GetOrigin()
-
 	// Root schema
 	if len(pes) == 0 || (len(pes) == 1 && pes[0] == "") {
 		err := s.db.View(func(txn *badger.Txn) error {
diff --git a/pkg/store/persiststore/persiststore_test.go b/pkg/store/persiststore/persiststore_test.go
index 6bc6ab5..c6d739e 100644
--- a/pkg/store/persiststore/persiststore_test.go
+++ b/pkg/store/persiststore/persiststore_test.go
@@ -18,8 +18,10 @@ import (
 	"context"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/dgraph-io/badger/v4"
+	"github.com/jellydator/ttlcache/v3"
 	"github.com/sdcio/schema-server/pkg/config"
 	"github.com/sdcio/schema-server/pkg/schema"
 	"github.com/sdcio/schema-server/pkg/store"
@@ -45,6 +47,18 @@ func newTestStore(t *testing.T) *persistStore {
 	return &persistStore{db: db}
 }
 
+func newTestStoreWithCache(t *testing.T) *persistStore {
+	t.Helper()
+	ps := newTestStore(t)
+	ps.cache = ttlcache.New[cacheKey, *sdcpb.GetSchemaResponse](
+		ttlcache.WithTTL[cacheKey, *sdcpb.GetSchemaResponse](time.Minute),
+		ttlcache.WithCapacity[cacheKey, *sdcpb.GetSchemaResponse](128),
+	)
+	go ps.cache.Start()
+	t.Cleanup(func() { ps.cache.Stop() })
+	return ps
+}
+
 func testSchemaKey() store.SchemaKey {
 	return store.SchemaKey{
 		Name:    "M",
@@ -282,3 +296,53 @@ func TestGetSchema_ModulePrefixedPathResolves(t *testing.T) {
 		t.Fatalf("unexpected schema name: %q", got)
 	}
 }
+
+// TestGetSchema_CacheKeyIncludesOrigin ensures the TTL cache distinguishes requests that
+// differ only by gNMI Path.origin (module hint); otherwise the second call would hit
+// the first resolution and return the wrong SchemaElem.
+func TestGetSchema_CacheKeyIncludesOrigin(t *testing.T) {
+	ps := newTestStoreWithCache(t)
+	sk := testSchemaKey()
+	insertSchemaMeta(t, ps, sk)
+	insertRootEntry(t, ps, sk, []string{"modA", "modB"})
+	insertEntry(t, ps, sk, []string{"modA", "dup"}, &sdcpb.SchemaElem{
+		Schema: &sdcpb.SchemaElem_Container{Container: &sdcpb.ContainerSchema{Name: "from-a"}},
+	})
+	insertEntry(t, ps, sk, []string{"modB", "dup"}, &sdcpb.SchemaElem{
+		Schema: &sdcpb.SchemaElem_Container{Container: &sdcpb.ContainerSchema{Name: "from-b"}},
+	})
+
+	path := &sdcpb.Path{Elem: []*sdcpb.PathElem{{Name: "dup"}}}
+	rspA, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Origin: "modA", Elem: path.Elem},
+	})
+	if err != nil {
+		t.Fatalf("GetSchema modA: %v", err)
+	}
+	if rspA.GetSchema().GetContainer().GetName() != "from-a" {
+		t.Fatalf("modA: got %q", rspA.GetSchema().GetContainer().GetName())
+	}
+
+	rspB, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Origin: "modB", Elem: path.Elem},
+	})
+	if err != nil {
+		t.Fatalf("GetSchema modB: %v", err)
+	}
+	if rspB.GetSchema().GetContainer().GetName() != "from-b" {
+		t.Fatalf("modB: got %q, want from-b (cache likely ignored origin)", rspB.GetSchema().GetContainer().GetName())
+	}
+
+	rspArev, err := ps.GetSchema(context.Background(), &sdcpb.GetSchemaRequest{
+		Schema: &sdcpb.Schema{Name: sk.Name, Vendor: sk.Vendor, Version: sk.Version},
+		Path:   &sdcpb.Path{Origin: "modA", Elem: path.Elem},
+	})
+	if err != nil {
+		t.Fatalf("GetSchema modA again: %v", err)
+	}
+	if rspArev.GetSchema().GetContainer().GetName() != "from-a" {
+		t.Fatalf("modA revisit: got %q", rspArev.GetSchema().GetContainer().GetName())
+	}
+}

From 3c08ab3c7070cc2b3fcbc123cda70aa98c942965 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 13:27:26 +0200
Subject: [PATCH 08/10] Align memstore and schema lookup with persiststore path
 semantics.

Apply Path.origin as a first-element module hint in memstore, accept
module-name prefixes at the root in FindPossibleModules, and add
parsePathElems plus memstore/persiststore GetSchema parity tests (with
separate compiled schemas so persiststore AddSchema Reset does not
invalidate memstore).
---
 pkg/schema/object.go                          |   2 +-
 pkg/store/memstore/memstore.go                |  17 +-
 .../persiststore/getschema_parity_test.go     | 159 ++++++++++++++++++
 .../testdata/origin-dup/moda.yang             |  12 ++
 .../testdata/origin-dup/modb.yang             |  12 ++
 pkg/store/persiststore/yang_helpers_test.go   |  50 ++++++
 6 files changed, 249 insertions(+), 3 deletions(-)
 create mode 100644 pkg/store/persiststore/getschema_parity_test.go
 create mode 100644 pkg/store/persiststore/testdata/origin-dup/moda.yang
 create mode 100644 pkg/store/persiststore/testdata/origin-dup/modb.yang
 create mode 100644 pkg/store/persiststore/yang_helpers_test.go

diff --git a/pkg/schema/object.go b/pkg/schema/object.go
index d5cb05f..2de6579 100644
--- a/pkg/schema/object.go
+++ b/pkg/schema/object.go
@@ -120,7 +120,7 @@ func (sc *Schema) FindPossibleModulesForPathElement(e *yang.Entry, pathElement s
 	case e.Node == nil:
 		entries := make([]*yang.Entry, 0)
 		for _, entry := range sc.root.Dir {
-			if ee, ok := entry.Dir[path]; ok && (!foundPrefix || ee.Prefix.Name == prefix) {
+			if ee, ok := entry.Dir[path]; ok && (!foundPrefix || ee.Prefix.Name == prefix || entry.Name == prefix) {
 				entries = append(entries, entry)
 			}
 		}
diff --git a/pkg/store/memstore/memstore.go b/pkg/store/memstore/memstore.go
index a3dd51b..fec7b18 100644
--- a/pkg/store/memstore/memstore.go
+++ b/pkg/store/memstore/memstore.go
@@ -17,6 +17,7 @@ package memstore
 import (
 	"context"
 	"sort"
+	"strings"
 	"sync"
 
 	"github.com/openconfig/goyang/pkg/yang"
@@ -43,6 +44,18 @@ func New() store.Store {
 	}
 }
 
+func pathStringsForGetSchema(req *sdcpb.GetSchemaRequest) []string {
+	pes := utils.ToStrings(req.GetPath(), false, true)
+	// Match persiststore: gNMI Path.origin acts as a module hint for the first path
+	// element when that element has no embedded "module:name" prefix.
+	if p := req.GetPath(); p != nil && p.GetOrigin() != "" && len(pes) > 0 {
+		if !strings.Contains(pes[0], ":") {
+			pes = append([]string{p.GetOrigin() + ":" + pes[0]}, pes[1:]...)
+		}
+	}
+	return pes
+}
+
 func (s *memStore) GetSchema(ctx context.Context, req *sdcpb.GetSchemaRequest) (*sdcpb.GetSchemaResponse, error) {
 	s.ms.RLock()
 	defer s.ms.RUnlock()
@@ -50,7 +63,7 @@ func (s *memStore) GetSchema(ctx context.Context, req *sdcpb.GetSchemaRequest) (
 	if reqSchema == nil {
 		return nil, status.Error(codes.InvalidArgument, "missing schema details")
 	}
-	pes := utils.ToStrings(req.GetPath(), false, true)
+	pes := pathStringsForGetSchema(req)
 
 	sc, ok := s.schemas[store.SchemaKey{Name: reqSchema.Name, Vendor: reqSchema.Vendor, Version: reqSchema.Version}]
 	if !ok {
@@ -309,7 +322,7 @@ func (s *memStore) GetSchemaElements(ctx context.Context, req *sdcpb.GetSchemaRe
 	if !ok {
 		return nil, status.Errorf(codes.InvalidArgument, "unknown schema %v", reqSchema)
 	}
-	pes := utils.ToStrings(req.GetPath(), false, true)
+	pes := pathStringsForGetSchema(req)
 
 	sch := make(chan *sdcpb.SchemaElem)
 	ych := make(chan *yang.Entry)
diff --git a/pkg/store/persiststore/getschema_parity_test.go b/pkg/store/persiststore/getschema_parity_test.go
new file mode 100644
index 0000000..ecab70a
--- /dev/null
+++ b/pkg/store/persiststore/getschema_parity_test.go
@@ -0,0 +1,159 @@
+// Copyright 2024 Nokia
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package persiststore
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"github.com/sdcio/schema-server/pkg/config"
+	"github.com/sdcio/schema-server/pkg/schema"
+	"github.com/sdcio/schema-server/pkg/store"
+	"github.com/sdcio/schema-server/pkg/store/memstore"
+	sdcpb "github.com/sdcio/sdc-protos/sdcpb"
+	"google.golang.org/protobuf/proto"
+)
+
+// assertGetSchemaParity checks that memstore and persiststore return the same outcome
+// for the same GetSchemaRequest (memstore is the reference implementation).
+func assertGetSchemaParity(t *testing.T, ms, ps store.Store, sk *sdcpb.Schema, req *sdcpb.GetSchemaRequest) {
+	t.Helper()
+	req = proto.Clone(req).(*sdcpb.GetSchemaRequest)
+	if req.Schema == nil {
+		req.Schema = sk
+	}
+	req.WithDescription = false
+
+	ctx := context.Background()
+	mRsp, mErr := ms.GetSchema(ctx, proto.Clone(req).(*sdcpb.GetSchemaRequest))
+	pRsp, pErr := ps.GetSchema(ctx, proto.Clone(req).(*sdcpb.GetSchemaRequest))
+
+	switch {
+	case mErr != nil && pErr != nil:
+		return
+	case mErr != nil:
+		t.Fatalf("memstore error %v, persiststore ok: %v", mErr, pRsp)
+	case pErr != nil:
+		t.Fatalf("persiststore error %v, memstore ok: %v", pErr, mRsp)
+	}
+
+	if !proto.Equal(mRsp.GetSchema(), pRsp.GetSchema()) {
+		t.Fatalf("SchemaElem mismatch.\nmemstore: %s\npersist: %s",
+			mRsp.GetSchema().String(), pRsp.GetSchema().String())
+	}
+}
+
+func newParityStores(t *testing.T, cfg *config.SchemaConfig) (ms store.Store, ps *persistStore, sk *sdcpb.Schema) {
+	t.Helper()
+	scMem, err := schema.NewSchema(cfg)
+	if err != nil {
+		t.Fatalf("NewSchema (memstore): %v", err)
+	}
+	scPersist, err := schema.NewSchema(cfg)
+	if err != nil {
+		t.Fatalf("NewSchema (persiststore): %v", err)
+	}
+	ms = memstore.New()
+	if err := ms.AddSchema(scMem); err != nil {
+		t.Fatalf("memstore AddSchema: %v", err)
+	}
+	ps = newTestStore(t)
+	if err := ps.AddSchema(scPersist); err != nil {
+		t.Fatalf("persiststore AddSchema: %v", err)
+	}
+	sk = &sdcpb.Schema{Name: scMem.Name(), Vendor: scMem.Vendor(), Version: scMem.Version()}
+	return ms, ps, sk
+}
+
+func TestGetSchema_MemstoreParity_AugmentFixture(t *testing.T) {
+	td := filepath.Join("testdata", "augment")
+	cfg := &config.SchemaConfig{
+		Name:    "parity-augment",
+		Vendor:  "test",
+		Version: "0",
+		Files:   []string{filepath.Join(td, "aug-base.yang"), filepath.Join(td, "aug-extra.yang")},
+	}
+	ms, ps, sk := newParityStores(t, cfg)
+
+	cases := []struct {
+		name string
+		path *sdcpb.Path
+	}{
+		{
+			name: "module_less_target_native_leaf",
+			path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+				{Name: "target"}, {Name: "native-leaf"},
+			}},
+		},
+		{
+			name: "first_elem_module_prefixed",
+			path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+				{Name: "aug-base:target"}, {Name: "native-leaf"},
+			}},
+		},
+		{
+			name: "augment_only_leaf_module_less",
+			path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+				{Name: "target"}, {Name: "augment-only-leaf"},
+			}},
+		},
+		{
+			name: "first_prefixed_inner_module_less",
+			path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+				{Name: "aug-base:target"}, {Name: "augment-only-leaf"},
+			}},
+		},
+		{
+			name: "inner_elem_has_module_prefix",
+			path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+				{Name: "aug-base:target"}, {Name: "aug-extra:augment-only-leaf"},
+			}},
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			assertGetSchemaParity(t, ms, ps, sk, &sdcpb.GetSchemaRequest{Schema: sk, Path: tc.path})
+		})
+	}
+}
+
+func TestGetSchema_MemstoreParity_OriginDisambiguatesDupTopLevel(t *testing.T) {
+	td := filepath.Join("testdata", "origin-dup")
+	cfg := &config.SchemaConfig{
+		Name:    "parity-origin-dup",
+		Vendor:  "test",
+		Version: "0",
+		Files:   []string{filepath.Join(td, "moda.yang"), filepath.Join(td, "modb.yang")},
+	}
+	ms, ps, sk := newParityStores(t, cfg)
+
+	basePath := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "dup"}, {Name: "marker"},
+	}}
+
+	t.Run("origin_moda", func(t *testing.T) {
+		assertGetSchemaParity(t, ms, ps, sk, &sdcpb.GetSchemaRequest{
+			Schema: sk,
+			Path:   &sdcpb.Path{Origin: "moda", Elem: basePath.Elem},
+		})
+	})
+	t.Run("origin_modb", func(t *testing.T) {
+		assertGetSchemaParity(t, ms, ps, sk, &sdcpb.GetSchemaRequest{
+			Schema: sk,
+			Path:   &sdcpb.Path{Origin: "modb", Elem: basePath.Elem},
+		})
+	})
+}
diff --git a/pkg/store/persiststore/testdata/origin-dup/moda.yang b/pkg/store/persiststore/testdata/origin-dup/moda.yang
new file mode 100644
index 0000000..200ae45
--- /dev/null
+++ b/pkg/store/persiststore/testdata/origin-dup/moda.yang
@@ -0,0 +1,12 @@
+module moda {
+  yang-version 1.1;
+  namespace "urn:schema-server:test:moda";
+  prefix ma;
+
+  container dup {
+    leaf marker {
+      type string;
+      default "from-moda";
+    }
+  }
+}
diff --git a/pkg/store/persiststore/testdata/origin-dup/modb.yang b/pkg/store/persiststore/testdata/origin-dup/modb.yang
new file mode 100644
index 0000000..0fc1a47
--- /dev/null
+++ b/pkg/store/persiststore/testdata/origin-dup/modb.yang
@@ -0,0 +1,12 @@
+module modb {
+  yang-version 1.1;
+  namespace "urn:schema-server:test:modb";
+  prefix mb;
+
+  container dup {
+    leaf marker {
+      type string;
+      default "from-modb";
+    }
+  }
+}
diff --git a/pkg/store/persiststore/yang_helpers_test.go b/pkg/store/persiststore/yang_helpers_test.go
new file mode 100644
index 0000000..14fdcb4
--- /dev/null
+++ b/pkg/store/persiststore/yang_helpers_test.go
@@ -0,0 +1,50 @@
+// Copyright 2024 Nokia
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package persiststore
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestParsePathElems(t *testing.T) {
+	t.Parallel()
+	cases := []struct {
+		in   []string
+		want []pathElem
+	}{
+		// nil input yields empty slice (not nil) from make(…, 0, len(pes))
+		{nil, []pathElem{}},
+		{[]string{}, []pathElem{}},
+		{[]string{"a"}, []pathElem{{name: "a"}}},
+		{[]string{"m:a"}, []pathElem{{module: "m", name: "a"}}},
+		{[]string{"foo:bar:baz"}, []pathElem{{module: "foo", name: "bar:baz"}}},
+		{[]string{":edge"}, []pathElem{{name: ":edge"}}},
+		{[]string{"no-colon-here"}, []pathElem{{name: "no-colon-here"}}},
+		{
+			[]string{"aug-base:target", "aug-extra:augment-only-leaf"},
+			[]pathElem{
+				{module: "aug-base", name: "target"},
+				{module: "aug-extra", name: "augment-only-leaf"},
+			},
+		},
+	}
+	for _, tc := range cases {
+		got := parsePathElems(tc.in)
+		if !reflect.DeepEqual(got, tc.want) {
+			t.Fatalf("parsePathElems(%q) = %#v, want %#v", tc.in, got, tc.want)
+		}
+	}
+}

From 87db0cfef8fb1ae24e50b519c5d328b874765d00 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Fri, 15 May 2026 13:40:01 +0200
Subject: [PATCH 09/10] fix(schema): resolve BuildPath first hop via Dir and
 Augmented

BuildPath used only e.Dir[pe[0]] after module resolution; align with
buildPath/getEntry by using entryChildByName and pathElemLocalName.
Adds a regression test for an augmented-only direct child under a module.
---
 pkg/schema/object.go      | 14 +++++++++----
 pkg/schema/object_test.go | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/pkg/schema/object.go b/pkg/schema/object.go
index 2de6579..1d26391 100644
--- a/pkg/schema/object.go
+++ b/pkg/schema/object.go
@@ -190,8 +190,7 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 
 			// prefix will be in [0] if exists, so path will always be in last index
 			// compare name without prefix
-			pathElements := strings.SplitN(pe[0], ":", 2)
-			if ee.Name != pathElements[len(pathElements)-1] {
+			if ee.Name != pathElemLocalName(pe[0]) {
 				continue
 			}
 			return getEntry(ee, pe[1:])
@@ -201,6 +200,13 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 	}
 }
 
+// pathElemLocalName returns the YANG identifier after an optional single "prefix:".
+// It matches the name comparison rule used in getEntry for prefixed path elements.
+func pathElemLocalName(s string) string {
+	parts := strings.SplitN(s, ":", 2)
+	return parts[len(parts)-1]
+}
+
 // entryChildByName returns a direct child of parent by name, checking Dir then Augmented.
 func entryChildByName(parent *yang.Entry, name string) *yang.Entry {
 	if parent == nil {
@@ -237,7 +243,7 @@ func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 		if e == nil {
 			return fmt.Errorf("module %q not found", first)
 		}
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pathElemLocalName(pe[0])); ee != nil {
 			err := sc.buildPath(pe, p, ee)
 			if err != nil {
 				return err
@@ -250,7 +256,7 @@ func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 	}
 	// try children
 	for _, e := range sc.root.Dir {
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pathElemLocalName(pe[0])); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 	}
diff --git a/pkg/schema/object_test.go b/pkg/schema/object_test.go
index df02ddf..a4d9ddb 100644
--- a/pkg/schema/object_test.go
+++ b/pkg/schema/object_test.go
@@ -18,6 +18,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"sort"
+	"sync"
 	"testing"
 
 	"github.com/openconfig/goyang/pkg/yang"
@@ -517,6 +518,46 @@ func TestSchema_BuildPath_AugmentedUnderCase(t *testing.T) {
 	}
 }
 
+// TestSchema_BuildPath_bootstrapAugmentedFirstHop verifies BuildPath resolves the
+// first in-module segment using the same Dir+Augmented lookup as buildPath
+// (regression: bootstrap used e.Dir[pe[0]] only).
+func TestSchema_BuildPath_bootstrapAugmentedFirstHop(t *testing.T) {
+	augLeaf := &yang.Entry{
+		Name:   "aug-only",
+		Kind:   yang.LeafEntry,
+		Type:   &yang.YangType{Kind: yang.Ystring},
+		Prefix: &yang.Value{Name: "pfx"},
+	}
+	mod := &yang.Entry{
+		Name:      "testmod",
+		Kind:      yang.DirectoryEntry,
+		Dir:       map[string]*yang.Entry{},
+		Augmented: []*yang.Entry{augLeaf},
+		Prefix:    &yang.Value{Name: "pfx"},
+	}
+	augLeaf.Parent = mod
+	root := &yang.Entry{
+		Name: RootName,
+		Kind: yang.DirectoryEntry,
+		Dir:  map[string]*yang.Entry{"testmod": mod},
+	}
+	sc := &Schema{
+		root:   root,
+		config: &config.SchemaConfig{Name: "t", Vendor: "t", Version: "0"},
+		m:      new(sync.RWMutex),
+	}
+	p := &sdcpb.Path{}
+	if err := sc.BuildPath([]string{"testmod:aug-only"}, p); err != nil {
+		t.Fatalf("BuildPath: %v", err)
+	}
+	want := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "testmod:aug-only"},
+	}}
+	if !comparePaths(p, want) {
+		t.Fatalf("got %v want %v", p, want)
+	}
+}
+
 func comparePathElem(pe1, pe2 *sdcpb.PathElem) bool {
 	if pe1 == nil {
 		return pe2 == nil

From 118e2eb2112430d712dbae04ae7062be2953e358 Mon Sep 17 00:00:00 2001
From: steiler <github@vahlenkamp.net>
Date: Mon, 18 May 2026 12:29:49 +0200
Subject: [PATCH 10/10] remove iteration of e.Augmented

---
 pkg/schema/container.go                       | 20 ---------
 pkg/schema/object.go                          | 45 +++----------------
 pkg/schema/object_test.go                     | 16 +++----
 pkg/schema/references.go                      | 15 -------
 pkg/schema/schema.go                          | 14 ------
 .../persiststore/getschema_parity_test.go     | 21 +++++++++
 pkg/store/persiststore/persiststore.go        | 34 ++++++++------
 .../testdata/module-root-parity/has-data.yang | 11 +++++
 .../module-root-parity/types-only.yang        |  9 ++++
 9 files changed, 75 insertions(+), 110 deletions(-)
 create mode 100644 pkg/store/persiststore/testdata/module-root-parity/has-data.yang
 create mode 100644 pkg/store/persiststore/testdata/module-root-parity/types-only.yang

diff --git a/pkg/schema/container.go b/pkg/schema/container.go
index 7d9d933..6342f3e 100644
--- a/pkg/schema/container.go
+++ b/pkg/schema/container.go
@@ -108,16 +108,6 @@ func getMandatoryChildren(e *yang.Entry) []*sdcpb.MandatoryChild {
 			result = append(result, c)
 		}
 	}
-	// include mandatory children from augments
-	for _, v := range e.Augmented {
-		if v.Mandatory == yang.TSTrue {
-			c := &sdcpb.MandatoryChild{
-				Name:    v.Name,
-				IsState: isState(v),
-			}
-			result = append(result, c)
-		}
-	}
 	return result
 }
 
@@ -167,16 +157,6 @@ func getChoiceInfo(e *yang.Entry) *sdcpb.ChoiceInfo {
 
 		processChoice(de, ci)
 	}
-	// also consider choices introduced via augments
-	for _, de := range e.Augmented {
-		if !de.IsChoice() {
-			continue
-		}
-		if ci == nil {
-			ci = &sdcpb.ChoiceInfo{Choice: map[string]*sdcpb.ChoiceInfoChoice{}}
-		}
-		processChoice(de, ci)
-	}
 	return ci
 }
 
diff --git a/pkg/schema/object.go b/pkg/schema/object.go
index 1d26391..1424f7c 100644
--- a/pkg/schema/object.go
+++ b/pkg/schema/object.go
@@ -182,7 +182,7 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 		}
 		return e, nil
 	default:
-		if e.Dir == nil && len(e.Augmented) == 0 {
+		if e.Dir == nil {
 			return nil, errors.New("not found")
 		}
 		for _, ee := range getChildren(e) {
@@ -207,20 +207,13 @@ func pathElemLocalName(s string) string {
 	return parts[len(parts)-1]
 }
 
-// entryChildByName returns a direct child of parent by name, checking Dir then Augmented.
+// entryChildByName returns a direct child of parent by name from Dir.
+// Goyang merges augment children into Dir; Augmented holds shallow augment-block copies only.
 func entryChildByName(parent *yang.Entry, name string) *yang.Entry {
 	if parent == nil {
 		return nil
 	}
-	if ce, ok := parent.Dir[name]; ok {
-		return ce
-	}
-	for _, ae := range parent.Augmented {
-		if ae != nil && ae.Name == name {
-			return ae
-		}
-	}
-	return nil
+	return parent.Dir[name]
 }
 
 func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
@@ -335,11 +328,6 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 					return err
 				}
 			}
-			for _, entry := range e.Augmented {
-				if handled, err := tryChoiceChild(e, entry); handled {
-					return err
-				}
-			}
 			return fmt.Errorf("choice %s - unknown element %s", e.Name, pe[0])
 		}
 	case e.IsCase():
@@ -397,7 +385,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 func getChildren(e *yang.Entry) []*yang.Entry {
 	switch {
 	case e.IsChoice(), e.IsCase(), e.IsContainer(), e.IsList():
-		rs := make([]*yang.Entry, 0, len(e.Dir)+len(e.Augmented))
+		rs := make([]*yang.Entry, 0, len(e.Dir))
 		for _, ee := range e.Dir {
 			if ee.IsChoice() || ee.IsCase() {
 				rs = append(rs, getChildren(ee)...)
@@ -405,14 +393,6 @@ func getChildren(e *yang.Entry) []*yang.Entry {
 			}
 			rs = append(rs, ee)
 		}
-		// add augmented children as well
-		for _, ee := range e.Augmented {
-			if ee.IsChoice() || ee.IsCase() {
-				rs = append(rs, getChildren(ee)...)
-				continue
-			}
-			rs = append(rs, ee)
-		}
 		//sort.Slice(rs, sortFn(rs))
 		return rs
 		// case e.IsCase():
@@ -618,18 +598,13 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 		return nil, fmt.Errorf("unknown element %s", pe[0])
 	}
 
-	// scan choice nodes in both direct and augmented children
+	// scan choice nodes under e (goyang merges augmented choices into Dir)
 	choices := make([]*yang.Entry, 0)
 	for _, child := range e.Dir {
 		if child != nil && child.IsChoice() {
 			choices = append(choices, child)
 		}
 	}
-	for _, child := range e.Augmented {
-		if child != nil && child.IsChoice() {
-			choices = append(choices, child)
-		}
-	}
 
 	for _, choice := range choices {
 		// implicit case: choice directly contains the data node
@@ -645,14 +620,6 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 				return target, nil
 			}
 		}
-		for _, cc := range choice.Augmented {
-			if cc == nil || !cc.IsCase() {
-				continue
-			}
-			if target := entryChildByName(cc, pe[1]); target != nil {
-				return target, nil
-			}
-		}
 	}
 	return nil, fmt.Errorf("unknown element %s", pe[1])
 }
diff --git a/pkg/schema/object_test.go b/pkg/schema/object_test.go
index a4d9ddb..3a0df42 100644
--- a/pkg/schema/object_test.go
+++ b/pkg/schema/object_test.go
@@ -518,10 +518,9 @@ func TestSchema_BuildPath_AugmentedUnderCase(t *testing.T) {
 	}
 }
 
-// TestSchema_BuildPath_bootstrapAugmentedFirstHop verifies BuildPath resolves the
-// first in-module segment using the same Dir+Augmented lookup as buildPath
-// (regression: bootstrap used e.Dir[pe[0]] only).
-func TestSchema_BuildPath_bootstrapAugmentedFirstHop(t *testing.T) {
+// TestSchema_BuildPath_bootstrapModuleChild verifies BuildPath resolves the first
+// in-module segment using Dir lookup on the module entry (goyang merges augments into Dir).
+func TestSchema_BuildPath_bootstrapModuleChild(t *testing.T) {
 	augLeaf := &yang.Entry{
 		Name:   "aug-only",
 		Kind:   yang.LeafEntry,
@@ -529,11 +528,10 @@ func TestSchema_BuildPath_bootstrapAugmentedFirstHop(t *testing.T) {
 		Prefix: &yang.Value{Name: "pfx"},
 	}
 	mod := &yang.Entry{
-		Name:      "testmod",
-		Kind:      yang.DirectoryEntry,
-		Dir:       map[string]*yang.Entry{},
-		Augmented: []*yang.Entry{augLeaf},
-		Prefix:    &yang.Value{Name: "pfx"},
+		Name:   "testmod",
+		Kind:   yang.DirectoryEntry,
+		Dir:    map[string]*yang.Entry{"aug-only": augLeaf},
+		Prefix: &yang.Value{Name: "pfx"},
 	}
 	augLeaf.Parent = mod
 	root := &yang.Entry{
diff --git a/pkg/schema/references.go b/pkg/schema/references.go
index 3a1e6d7..a3f646a 100644
--- a/pkg/schema/references.go
+++ b/pkg/schema/references.go
@@ -70,21 +70,6 @@ func (sc *Schema) buildReferences(e *yang.Entry) error {
 			return err
 		}
 	}
-	// also recurse into augmented children
-	for _, ce := range e.Augmented {
-		if ce.IsCase() || ce.IsChoice() {
-			for _, cce := range ce.Dir {
-				err := sc.buildReferences(cce)
-				if err != nil {
-					return err
-				}
-			}
-		}
-		err := sc.buildReferences(ce)
-		if err != nil {
-			return err
-		}
-	}
 	return nil
 }
 
diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
index dbeea39..08bc4d5 100644
--- a/pkg/schema/schema.go
+++ b/pkg/schema/schema.go
@@ -141,13 +141,6 @@ func (s *Schema) Walk(e *yang.Entry, fn func(ec *yang.Entry) error) error {
 				return err
 			}
 		}
-		// also walk augments merged into this entry
-		for _, ee := range e.Augmented {
-			err = s.Walk(ee, fn)
-			if err != nil {
-				return err
-			}
-		}
 		return nil
 	}
 	err = fn(e)
@@ -160,13 +153,6 @@ func (s *Schema) Walk(e *yang.Entry, fn func(ec *yang.Entry) error) error {
 			return err
 		}
 	}
-	// include augmented children at every level
-	for _, ce := range e.Augmented {
-		err = s.Walk(ce, fn)
-		if err != nil {
-			return err
-		}
-	}
 	return nil
 }
 
diff --git a/pkg/store/persiststore/getschema_parity_test.go b/pkg/store/persiststore/getschema_parity_test.go
index ecab70a..95e0cdf 100644
--- a/pkg/store/persiststore/getschema_parity_test.go
+++ b/pkg/store/persiststore/getschema_parity_test.go
@@ -130,6 +130,27 @@ func TestGetSchema_MemstoreParity_AugmentFixture(t *testing.T) {
 	}
 }
 
+func TestGetSchema_MemstoreParity_SingleSegmentModuleRoot_TypesOnlyModule(t *testing.T) {
+	td := filepath.Join("testdata", "module-root-parity")
+	cfg := &config.SchemaConfig{
+		Name:    "parity-module-root",
+		Vendor:  "test",
+		Version: "0",
+		Files: []string{
+			filepath.Join(td, "types-only.yang"),
+			filepath.Join(td, "has-data.yang"),
+		},
+	}
+	ms, ps, sk := newParityStores(t, cfg)
+
+	assertGetSchemaParity(t, ms, ps, sk, &sdcpb.GetSchemaRequest{
+		Schema: sk,
+		Path: &sdcpb.Path{Elem: []*sdcpb.PathElem{
+			{Name: "parity-types-only"},
+		}},
+	})
+}
+
 func TestGetSchema_MemstoreParity_OriginDisambiguatesDupTopLevel(t *testing.T) {
 	td := filepath.Join("testdata", "origin-dup")
 	cfg := &config.SchemaConfig{
diff --git a/pkg/store/persiststore/persiststore.go b/pkg/store/persiststore/persiststore.go
index 0b7c0b4..45e0f21 100644
--- a/pkg/store/persiststore/persiststore.go
+++ b/pkg/store/persiststore/persiststore.go
@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"slices"
 	"sort"
 	"strings"
 	"time"
@@ -578,12 +579,6 @@ func (s *persistStore) addSchemaElem(wb *badger.WriteBatch, sc *schema.Schema, e
 				return err
 			}
 		}
-		for _, ee := range e.Augmented {
-			err := s.addSchemaElem(wb, sc, ee)
-			if err != nil {
-				return err
-			}
-		}
 		return nil
 	}
 	// build entry key
@@ -602,19 +597,13 @@ func (s *persistStore) addSchemaElem(wb *badger.WriteBatch, sc *schema.Schema, e
 	if err != nil {
 		return err
 	}
-	// do the same for entry children (native definitions and augment-only nodes)
+	// recurse into merged schema children (goyang places augment nodes in Dir)
 	for _, ee := range e.Dir {
 		err = s.addSchemaElem(wb, sc, ee)
 		if err != nil {
 			return err
 		}
 	}
-	for _, ee := range e.Augmented {
-		err = s.addSchemaElem(wb, sc, ee)
-		if err != nil {
-			return err
-		}
-	}
 	return nil
 }
 
@@ -798,6 +787,9 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 
 	// Parse per-element prefixes and build unprefixed names
 	pp := parsePathElems(pes)
+	// If the first element is later scoped by origin or "mod:" prefix, do not fall back to
+	// module-root keys (would return the wrong module after a failed scoped lookup).
+	firstPathElemUnscoped := origin == "" && (len(pp) == 0 || pp[0].module == "")
 	names := make([]string, 0, len(pp))
 	for _, pe := range pp {
 		names = append(names, pe.name)
@@ -860,6 +852,22 @@ func (s *persistStore) getSchema(_ context.Context, req *sdcpb.GetSchemaRequest,
 			}
 			return nil
 		}
+		// Single-segment path that is a compiled module name resolves to the module root
+		// entry (key buildEntryKey(sck, []string{module})), not [module, module] which the
+		// loop above probes. Memstore does this via Schema.GetEntry.
+		if firstPathElemUnscoped && len(names) == 1 && slices.Contains(allModules, names[0]) {
+			k := buildEntryKey(sck, []string{names[0]})
+			if item, err := txn.Get(k); err == nil && item != nil {
+				v, err := item.ValueCopy(nil)
+				if err != nil {
+					return err
+				}
+				if err := proto.Unmarshal(v, sce); err != nil {
+					return err
+				}
+				return nil
+			}
+		}
 		return fmt.Errorf("schema path not found: %s", req.GetPath())
 	})
 	if err != nil {
diff --git a/pkg/store/persiststore/testdata/module-root-parity/has-data.yang b/pkg/store/persiststore/testdata/module-root-parity/has-data.yang
new file mode 100644
index 0000000..b0ba3f0
--- /dev/null
+++ b/pkg/store/persiststore/testdata/module-root-parity/has-data.yang
@@ -0,0 +1,11 @@
+module parity-has-data {
+  yang-version 1.1;
+  namespace "urn:test:parity:data";
+  prefix pdata;
+
+  container rootbox {
+    leaf a {
+      type string;
+    }
+  }
+}
diff --git a/pkg/store/persiststore/testdata/module-root-parity/types-only.yang b/pkg/store/persiststore/testdata/module-root-parity/types-only.yang
new file mode 100644
index 0000000..d4c0618
--- /dev/null
+++ b/pkg/store/persiststore/testdata/module-root-parity/types-only.yang
@@ -0,0 +1,9 @@
+module parity-types-only {
+  yang-version 1.1;
+  namespace "urn:test:parity:typesonly";
+  prefix ptypes;
+
+  typedef example-string {
+    type string;
+  }
+}