ast+rego+topdown: external rule source support (#8600)

External rule sources let wrapping projects inject rules at evaluation
time instead of compile time. The compiler marks external packages in
the rule tree but doesn't index them. When topdown hits an external
node, it calls Lookup to get rules, compiles them on the fly with a
scoped compiler, grafts the result into the tree, and caches it for the
duration of the evaluation.

Sources can be isolated (default, no access to surrounding policy) or
non-isolated (can reference static rules and other external sources).
The ExternalRuleIndexCloser interface handles cleanup after evaluation.
Precompiled rules can skip compiler stages via SkippedStages to avoid
redundant work.

This includes:
* hooks: add BundlePreActivate hook This one is handy when registering
external sources.

* topdown: catch `ir == nil` rule index result
This wouldn't ordinarily happen: the compiler is checking refs before.
But in our use case, the SP rules may be configured to be able to reach
into the surrounding Rego (non-isolated mode). If that happens, the IR
lookup may indeed end up as `nil, nil`.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

internal/runtime/init/init.go

@@ -31,6 +31,7 @@ type InsertAndCompileOptions struct {
 	EnablePrintStatements bool
 	ParserOptions         ast.ParserOptions
 	BundleActivatorPlugin string
+	ExternalSources       *util.HasherMap[ast.Ref, ast.ExternalRuleSource]
 }
 
 // InsertAndCompileResult contains the output of the operation.
@@ -59,18 +60,31 @@ func InsertAndCompile(ctx context.Context, opts InsertAndCompileOptions) (*Inser
 		SetErrorLimit(opts.MaxErrors).
 		WithPathConflictsCheck(storage.NonEmpty(ctx, opts.Store, opts.Txn)).
 		WithEnablePrintStatements(opts.EnablePrintStatements)
+
+	// Apply external sources to the compiler before bundle activation.
+	// Bundle activation applies them again via compileModules, but we need them
+	// here too: there may be no bundles, or a custom activator plugin may not
+	// call compileModules.
+	if opts.ExternalSources != nil {
+		opts.ExternalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool {
+			compiler = compiler.WithExternalSource(ref, source)
+			return false
+		})
+	}
+
 	m := metrics.New()
 
 	activation := &bundle.ActivateOpts{
-		Ctx:           ctx,
-		Store:         opts.Store,
-		Txn:           opts.Txn,
-		Compiler:      compiler,
-		Metrics:       m,
-		Bundles:       opts.Bundles,
-		ExtraModules:  policies,
-		ParserOptions: opts.ParserOptions,
-		Plugin:        opts.BundleActivatorPlugin,
+		Ctx:             ctx,
+		Store:           opts.Store,
+		Txn:             opts.Txn,
+		Compiler:        compiler,
+		Metrics:         m,
+		Bundles:         opts.Bundles,
+		ExtraModules:    policies,
+		ExternalSources: opts.ExternalSources,
+		ParserOptions:   opts.ParserOptions,
+		Plugin:          opts.BundleActivatorPlugin,
 	}
 
 	err := bundle.Activate(activation)

v1/ast/compile.go

@@ -5,6 +5,7 @@
 package ast
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -128,6 +129,7 @@ type Compiler struct {
 
 	localvargen                *localVarGenerator
 	moduleLoader               ModuleLoader
+	externalSources            *util.HasherMap[Ref, ExternalRuleSource]
 	stages                     []stage
 	maxErrs                    int
 	errCount                   uint32
@@ -135,6 +137,7 @@ type Compiler struct {
 	sorted                     []string    // list of sorted module names
 	pathExists                 func([]string) (bool, error)
 	pathConflictCheckRoots     []string
+	injectedVirtual            func(Ref) bool // optional custom virtual document checker
 	after                      map[string][]CompilerStageDefinition
 	metrics                    metrics.Metrics
 	capabilities               *Capabilities                 // user-supplied capabilities
@@ -419,6 +422,7 @@ func NewCompiler() *Compiler {
 		Modules:               map[string]*Module{},
 		RewrittenVars:         map[Var]Var{},
 		Required:              &Capabilities{},
+		externalSources:       util.NewHasherMap[Ref, ExternalRuleSource](RefEqual),
 		maxErrs:               CompileErrorLimitDefault,
 		mu:                    &sync.Mutex{},
 		after:                 map[string][]CompilerStageDefinition{},
@@ -651,6 +655,14 @@ func (c *Compiler) QueryCompiler() QueryCompiler {
 	return newQueryCompiler(&c0)
 }
 
+// WithVirtual sets a custom virtual document checker on the compiler.
+// The provided function will be called during rule index building to determine
+// if additional refs should be considered virtual documents.
+func (c *Compiler) WithVirtual(fn func(Ref) bool) *Compiler {
+	c.injectedVirtual = fn
+	return c
+}
+
 // Compile runs the compilation process on the input modules. The compiled
 // version of the modules and associated data structures are stored on the
 // compiler. If the compilation process fails for any reason, the compiler will
@@ -894,9 +906,8 @@ func (c *Compiler) GetRulesDynamic(ref Ref) []*Rule {
 // Without the options, it would be excluded.
 func (c *Compiler) GetRulesDynamicWithOpts(ref Ref, opts RulesOptions) []*Rule {
 	node := c.RuleTree
-
 	set := map[*Rule]struct{}{}
-	var walk func(node *TreeNode, i int)
+	var walk func(*TreeNode, int)
 	walk = func(node *TreeNode, i int) {
 		switch {
 		case i >= len(ref):
@@ -1041,6 +1052,19 @@ func (c *Compiler) WithModuleLoader(f ModuleLoader) *Compiler {
 	return c
 }
 
+// WithExternalSource registers an external rule source for the given package
+// reference. When rules under this package are queried via RuleIndex, the
+// external source will be invoked to fetch all rules for the package. The
+// fetched rules are cached so the external source is only called once per
+// package.
+//
+// The package reference should be a fully qualified path (e.g., data.foo.bar).
+// All rule queries under this package will be handled by the external source.
+func (c *Compiler) WithExternalSource(packageRef Ref, source ExternalRuleSource) *Compiler {
+	c.externalSources.Put(packageRef, source)
+	return c
+}
+
 // WithDefaultRegoVersion sets the default Rego version to use when a module doesn't specify one;
 // such as when it's hand-crafted instead of parsed.
 func (c *Compiler) WithDefaultRegoVersion(regoVersion RegoVersion) *Compiler {
@@ -1113,10 +1137,14 @@ func (c *Compiler) counterAdd(name string, n uint64) {
 func (c *Compiler) buildRuleIndices() {
 
 	c.RuleTree.DepthFirst(func(node *TreeNode) bool {
-		if len(node.Values) == 0 {
+		if len(node.Values) == 0 && node.External == nil {
 			return false
 		}
-		rules := node.Values
+		if node.External != nil {
+			// Skip external sources - they build indices dynamically
+			return true
+		}
+		rules := node.Values // must be len > 0 here
 		hasNonGroundRef := false
 		for _, r := range rules {
 			hasNonGroundRef = !r.Head.Ref().IsGround()
@@ -1141,9 +1169,7 @@ func (c *Compiler) buildRuleIndices() {
 			}
 		}
 
-		index := newBaseDocEqIndex(func(ref Ref) bool {
-			return isVirtual(c.RuleTree, ref.GroundPrefix())
-		})
+		index := newBaseDocEqIndex(c.isVirtual)
 		if index.Build(rules) {
 			node.Index = index
 		}
@@ -1302,15 +1328,19 @@ func (c *Compiler) checkRuleConflicts() {
 			return false // go deeper
 		}
 
-		kinds := make(map[RuleKind]struct{}, len(node.Values))
+		rules := node.Values
+		if len(rules) == 0 {
+			return true // ?? right
+		}
+		kinds := make(map[RuleKind]struct{}, len(rules))
 		completeRules := 0
 		partialRules := 0
-		arities := make(map[int]struct{}, len(node.Values))
+		arities := make(map[int]struct{}, len(rules))
 		name := ""
 		var conflicts []Ref
 		defaultRules := make([]*Rule, 0)
 
-		for _, rule := range node.Values {
+		for _, rule := range rules {
 			r := rule
 			ref := r.Ref()
 			name = rw(ref.CopyNonGround()).String() // varRewriter operates in-place
@@ -1383,10 +1413,10 @@ func (c *Compiler) checkRuleConflicts() {
 
 		switch {
 		case conflicts != nil:
-			return !c.err(NewError(TypeErr, node.Values[0].Loc(), "rule %v conflicts with %v", name, conflicts))
+			return !c.err(NewError(TypeErr, rules[0].Loc(), "rule %v conflicts with %v", name, conflicts))
 
 		case len(kinds) > 1 || len(arities) > 1 || (completeRules >= 1 && partialRules >= 1):
-			return !c.err(NewError(TypeErr, node.Values[0].Loc(), "conflicting rules %v found", name))
+			return !c.err(NewError(TypeErr, rules[0].Loc(), "conflicting rules %v found", name))
 
 		case len(defaultRules) > 1:
 			buf := append(append(append(make([]byte, 0, 64), "multiple default rules "...), name...), " found at "...)
@@ -3455,6 +3485,17 @@ func (c *Compiler) setModuleTree() {
 
 func (c *Compiler) setRuleTree() {
 	c.RuleTree = NewRuleTree(c.ModuleTree)
+
+	// Add tree nodes for external source paths so evaluation knows to look there
+	c.externalSources.Iter(func(pkgRef Ref, source ExternalRuleSource) bool {
+		ri, err := source.Init(context.TODO(), pkgRef)
+		if err != nil {
+			c.err(NewError(CompileErr, nil, "failed to initialize external rule source for ref %v: %v", pkgRef, err))
+			return true
+		}
+		c.RuleTree.add(pkgRef, ri)
+		return false
+	})
 }
 
 func (c *Compiler) setGraph() {
@@ -4123,6 +4164,7 @@ func (n *ModuleTreeNode) DepthFirst(f func(*ModuleTreeNode) bool) {
 // rule path.
 type TreeNode struct {
 	Key      Value
+	External *ExternalIndex
 	Values   []*Rule
 	Children map[Value]*TreeNode
 	Sorted   []Value
@@ -4167,20 +4209,104 @@ func NewRuleTree(mtree *ModuleTreeNode) *TreeNode {
 	return &root
 }
 
-func (n *TreeNode) add(path Ref, rule *Rule) {
+func (n *TreeNode) add(path Ref, val any) {
 	node, tail := n.find(path)
 	if len(tail) > 0 {
-		sub := treeNodeFromRef(tail, rule)
+		sub := treeNodeFromRef(path, tail, val)
 		if node.Children == nil {
 			node.Children = make(map[Value]*TreeNode, 1)
 		}
 		node.Children[sub.Key] = sub
 		node.Sorted = append(node.Sorted, sub.Key)
-	} else if rule != nil {
-		node.Values = append(node.Values, rule)
+	} else if val != nil {
+		switch val := val.(type) {
+		case *Rule:
+			node.Values = append(node.Values, val)
+		case ExternalRuleIndex:
+			node.External = &ExternalIndex{
+				Index: val,
+				Ref:   path,
+			}
+		}
 	}
 }
 
+type ExternalIndex struct {
+	Index ExternalRuleIndex
+	Ref   Ref
+}
+
+func (ei *ExternalIndex) Tree(ctx context.Context, rt *TreeNode, prefix Ref, input *Term, m metrics.Metrics, reqMD map[string]any, respMD map[string]any) (*TreeNode, ExternalRuleIndex, error) {
+	resolver := &termResolver{input: input}
+
+	rules, updatedIndex, err := ei.Index.Lookup(ctx,
+		LookupResolver(resolver),
+		LookupMetrics(m),
+		LookupRequestMetadata(reqMD),
+		LookupResponseMetadata(respMD),
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+	c0 := NewCompiler()
+
+	if o := ei.Index.Opts(); o != nil {
+		if len(o.SkippedStages) > 0 {
+			c0.WithSkipStages(o.SkippedStages...)
+		}
+
+		if len(o.VisibleRefs) > 0 {
+			visible := o.VisibleRefs
+			c0.WithVirtual(func(ref Ref) bool {
+				return slices.ContainsFunc(visible, ref.HasPrefix) && rt.isVirtual(ref)
+			})
+		}
+	}
+
+	modules := make(map[string]*Module)
+	for _, rule := range rules {
+		pkgPathStr := rule.Module.Package.Path.String()
+		if mod, exists := modules[pkgPathStr]; exists {
+			mod.Rules = append(mod.Rules, rule)
+		} else {
+			modules[pkgPathStr] = &Module{
+				Package: &Package{Path: rule.Module.Package.Path},
+				Rules:   []*Rule{rule},
+			}
+		}
+	}
+	if m != nil {
+		t := m.Timer("external_lookup_compile_module")
+		t.Start()
+		defer t.Stop()
+	}
+	c0.Compile(modules)
+	if c0.Failed() {
+		return nil, nil, c0.Errors
+	}
+
+	node := c0.RuleTree.Find(prefix)
+	return node, updatedIndex, nil
+}
+
+type termResolver struct {
+	input *Term
+}
+
+func (r *termResolver) Resolve(ref Ref) (Value, error) {
+	if ref.HasPrefix(InputRootRef) {
+		if r.input == nil {
+			return nil, UnknownValueErr{}
+		}
+		v, err := r.input.Value.Find(ref[1:])
+		if err != nil {
+			return nil, UnknownValueErr{}
+		}
+		return v, nil
+	}
+	return nil, UnknownValueErr{}
+}
+
 // Size returns the number of rules in the tree.
 func (n *TreeNode) Size() (s int) {
 	for _, c := range n.Children {
@@ -4240,28 +4366,67 @@ func (n *TreeNode) DepthFirst(f func(*TreeNode) bool) {
 	}
 }
 
-func treeNodeFromRef(ref Ref, rule *Rule) *TreeNode {
-	depth := len(ref) - 1
-	key := ref[depth].Value
-	node := &TreeNode{
-		Key:      key,
-		Children: nil,
+func (c *Compiler) isVirtual(ref Ref) bool {
+	return (c.injectedVirtual != nil && c.injectedVirtual(ref)) ||
+		c.RuleTree.isVirtual(ref.GroundPrefix())
+}
+
+// isVirtual returns true if the ref is virtual (has rules).
+func (n *TreeNode) isVirtual(ref Ref) bool {
+	node := n
+	for i := range ref {
+		child := node.Child(ref[i].Value)
+		if child == nil {
+			return false
+		} else if len(child.Values) > 0 || child.External != nil {
+			return true
+		}
+		node = child
 	}
-	if rule != nil {
-		node.Values = []*Rule{rule}
+	return true
+}
+
+func treeNodeFromRef(ref, tail Ref, val any) *TreeNode {
+	if len(tail) == 0 {
+		node := &TreeNode{
+			Children: make(map[Value]*TreeNode),
+		}
+		attachValueToNode(node, ref, val)
+		return node
 	}
 
-	for i := len(ref) - 2; i >= 0; i-- {
-		key := ref[i].Value
+	depth := len(tail) - 1
+	node := &TreeNode{
+		Key: tail[depth].Value,
+	}
+	attachValueToNode(node, ref, val)
+
+	for i := depth - 1; i >= 0; i-- {
+		childKey := tail[i+1].Value
 		node = &TreeNode{
-			Key:      key,
-			Children: map[Value]*TreeNode{ref[i+1].Value: node},
-			Sorted:   []Value{ref[i+1].Value},
+			Key:      tail[i].Value,
+			Children: map[Value]*TreeNode{childKey: node},
+			Sorted:   []Value{childKey},
 		}
 	}
 	return node
 }
 
+func attachValueToNode(node *TreeNode, ref Ref, val any) {
+	if val == nil {
+		return
+	}
+	switch val := val.(type) {
+	case *Rule:
+		node.Values = append(node.Values, val)
+	case ExternalRuleIndex:
+		node.External = &ExternalIndex{
+			Index: val,
+			Ref:   ref,
+		}
+	}
+}
+
 // flattenChildren flattens all children's rule refs into a sorted array.
 func (n *TreeNode) flattenChildren() []Ref {
 	return n.flattenMatchingChildren(func(_ *Rule) bool { return true })
@@ -4288,6 +4453,60 @@ func (n *TreeNode) flattenMatchingChildren(f func(*Rule) bool) []Ref {
 	return util.SortedFunc(ret.s, RefCompare)
 }
 
+// Copy creates a shallow copy of the TreeNode suitable for augmentation.
+// Children map is copied recursively. Values slices are initially shared but
+// reallocated on modification (e.g., by MergeChild's append operation).
+func (n *TreeNode) Copy() *TreeNode {
+	if n == nil {
+		return nil
+	}
+
+	result := &TreeNode{
+		Key:      n.Key,
+		External: n.External,
+		Values:   n.Values,
+		Hide:     n.Hide,
+		Index:    n.Index,
+	}
+
+	if n.Children != nil {
+		result.Children = make(map[Value]*TreeNode, len(n.Children))
+		for k, v := range n.Children {
+			result.Children[k] = v.Copy()
+		}
+	}
+
+	if n.Sorted != nil {
+		result.Sorted = make([]Value, len(n.Sorted))
+		copy(result.Sorted, n.Sorted)
+	}
+
+	return result
+}
+
+// MergeChild merges another TreeNode into this node's children.
+func (n *TreeNode) MergeChild(key Value, other *TreeNode) {
+	if other == nil {
+		return
+	}
+
+	existing := n.Child(key)
+	if existing == nil {
+		if n.Children == nil {
+			n.Children = make(map[Value]*TreeNode)
+		}
+		n.Children[key] = other
+		n.Sorted = append(n.Sorted, key)
+		return
+	}
+
+	existing.Values = append(existing.Values, other.Values...)
+
+	for childKey, childNode := range other.Children {
+		existing.MergeChild(childKey, childNode)
+	}
+}
+
 // Graph represents the graph of dependencies between rules.
 type Graph struct {
 	adj    map[util.T]map[util.T]struct{}
@@ -4299,7 +4518,6 @@ type Graph struct {
 // NewGraph returns a new Graph based on modules. The list function must return
 // the rules referred to directly by the ref.
 func NewGraph(modules map[string]*Module, list func(Ref) []*Rule) *Graph {
-
 	graph := &Graph{
 		adj:    map[util.T]map[util.T]struct{}{},
 		radj:   map[util.T]map[util.T]struct{}{},
@@ -6676,8 +6894,8 @@ func validateWith(c *Compiler, unsafeBuiltinsMap map[string]struct{}, expr *Expr
 			// target is a function. It's probably wrong for arity-0 functions, but those are
 			// and edge case anyways.
 			if child := targetNode.Child(ref[len(ref)-1].Value); child != nil {
-				for _, v := range child.Values {
-					if len(v.Head.Args) > 0 {
+				for _, r := range child.Values {
+					if len(r.Head.Args) > 0 {
 						if ok, err := validateWithFunctionValue(c.builtins, unsafeBuiltinsMap, c.RuleTree, value); err != nil || ok {
 							return false, err // err may be nil
 						}
@@ -6690,8 +6908,8 @@ func validateWith(c *Compiler, unsafeBuiltinsMap map[string]struct{}, expr *Expr
 		if r, ok := value.Value.(Ref); ok {
 			// TODO: check that target ref doesn't exist?
 			if valueNode := c.RuleTree.Find(r); valueNode != nil {
-				for _, v := range valueNode.Values {
-					if len(v.Head.Args) > 0 {
+				for _, r := range valueNode.Values {
+					if len(r.Head.Args) > 0 {
 						return false, nil
 					}
 				}
@@ -6776,19 +6994,6 @@ func isBuiltinRefOrVar(bs map[string]*Builtin, unsafeBuiltinsMap map[string]stru
 	return false, nil
 }
 
-func isVirtual(node *TreeNode, ref Ref) bool {
-	for i := range ref {
-		child := node.Child(ref[i].Value)
-		if child == nil {
-			return false
-		} else if len(child.Values) > 0 {
-			return true
-		}
-		node = child
-	}
-	return true
-}
-
 func safetyErrorSlice(unsafe unsafeVars, rewritten map[Var]Var) (result Errors) {
 	if len(unsafe) == 0 {
 		return

v1/ast/compile_test.go

@@ -819,14 +819,15 @@ func TestRuleIndices(t *testing.T) {
 				c.sorted = append(c.sorted, strconv.Itoa(i))
 			}
 			compileStages(c, StageBuildRuleIndices)
+			t.Log(c.RuleTree.Dump())
 
 			for k, expIndex := range tc.exp {
 				kref := MustParseRef(k)
-				node := c.RuleTree.Find(kref)
-				if node == nil || node.Index == nil {
+				i := c.RuleIndex(kref)
+				if i == nil {
 					t.Fatalf("expected rule indices for %v", k)
 				}
-				index := node.Index.(*baseDocEqIndex)
+				index := i.(*baseDocEqIndex)
 				for _, expRef := range expIndex {
 					found := false
 					for _, r := range index.root.rules {
@@ -1199,7 +1200,7 @@ func TestRuleTree(t *testing.T) {
 		t.Fatalf("Expected user.system node to be visible")
 	}
 
-	if !isVirtual(tree, MustParseRef("data.a.b.empty")) {
+	if !tree.isVirtual(MustParseRef("data.a.b.empty")) {
 		t.Fatal("Expected data.a.b.empty to be virtual")
 	}
 

v1/ast/external_source.go

@@ -0,0 +1,128 @@
+// Copyright 2026 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package ast
+
+import (
+	"context"
+
+	"github.com/open-policy-agent/opa/v1/metrics"
+)
+
+type ExternalRuleSource interface {
+	// Refs returns the package refs that this source provides rules for.
+	// A source can provide rules for multiple packages.
+	Refs() []Ref
+
+	// Init returns an initialized [ExternalRuleIndex]. A `Ref` is provided
+	// so we know which package we're preparing if multiple Refs are external.
+	Init(context.Context, Ref) (ExternalRuleIndex, error)
+}
+
+// ExternalRuleIndex mirrors RuleIndex.Lookup(), but add a [context.Context] parameter.
+type ExternalRuleIndex interface {
+	// Opts returns the options for the ExternalRuleIndex. Returns nil if no
+	// options are configured.
+	Opts() *ExternalSourceOptions
+
+	// Lookup returns rules and optionally an updated ExternalRuleIndex instance.
+	// The returned ExternalRuleIndex (if non-nil) will be used for subsequent
+	// Lookup calls within the same evaluation context, allowing plugins to
+	// maintain per-evaluation state.
+	//
+	// Plugins can use two strategies:
+	// 1. Immutable: Return a new ExternalRuleIndex instance with updated state
+	// 2. Mutable: Update internal state and return self
+	//
+	// If the plugin does not need per-evaluation state, it can return nil for
+	// the ExternalRuleIndex, and the original instance will continue to be used.
+	Lookup(context.Context, ...LookupOption) ([]*Rule, ExternalRuleIndex, error)
+}
+
+// ExternalRuleIndexCloser is an optional interface for resource cleanup.
+type ExternalRuleIndexCloser interface {
+	ExternalRuleIndex
+	Close() error
+}
+
+// ExternalSourceOptions contains options for registering an external rule source.
+type ExternalSourceOptions struct {
+	// VisibleRefs controls which parts of the surrounding rule tree the external
+	// source can reference during compilation. By default (nil), the source is
+	// fully isolated and cannot access any surrounding policy. An empty slice
+	// is equivalent to nil (fully isolated).
+	//
+	// To allow access to the entire rule tree, use []Ref{MustParseRef("data")}.
+	// To allow access to specific subtrees only, list them explicitly, e.g.
+	// []Ref{MustParseRef("data.helpers")}. The external source can then
+	// reference rules under those prefixes but nothing else.
+	VisibleRefs []Ref
+
+	// SkippedStages allows external sources to skip stages in the dynamic compiler
+	// used with the externally-provided Rego. If, for example, the `[]*Rule` returned
+	// has already been compiled, we can skip all stages.
+	//
+	// For pre-compiled rules, prefer starting from AllStages() and removing only
+	// the stages you need (e.g. SetModuleTree, SetRuleTree, BuildRuleIndices).
+	// This is forward-compatible: new compiler stages added in future releases
+	// will be skipped automatically rather than running unexpectedly.
+	SkippedStages []StageID
+}
+
+// LookupOption is a functional option for ExternalRuleIndex.Lookup calls.
+type LookupOption func(*LookupOptions)
+
+// LookupOptions contains options for ExternalRuleIndex.Lookup calls.
+type LookupOptions struct {
+	metrics          metrics.Metrics
+	resolver         ValueResolver
+	requestMetadata  map[string]any
+	responseMetadata map[string]any
+}
+
+// Metrics returns the metrics instance from the options, or nil if not set.
+func (o *LookupOptions) Metrics() metrics.Metrics {
+	if o == nil {
+		return nil
+	}
+	return o.metrics
+}
+
+func (o *LookupOptions) Resolver() ValueResolver {
+	return o.resolver
+}
+
+func (o *LookupOptions) RequestMetadata() map[string]any {
+	return o.requestMetadata
+}
+
+func (o *LookupOptions) ResponseMetadata() map[string]any {
+	return o.responseMetadata
+}
+
+// LookupMetrics returns a LookupOption that sets the metrics instance
+// for the Lookup call.
+func LookupMetrics(m metrics.Metrics) LookupOption {
+	return func(opts *LookupOptions) {
+		opts.metrics = m
+	}
+}
+
+func LookupResolver(r ValueResolver) LookupOption {
+	return func(opts *LookupOptions) {
+		opts.resolver = r
+	}
+}
+
+func LookupRequestMetadata(m map[string]any) LookupOption {
+	return func(opts *LookupOptions) {
+		opts.requestMetadata = m
+	}
+}
+
+func LookupResponseMetadata(m map[string]any) LookupOption {
+	return func(opts *LookupOptions) {
+		opts.responseMetadata = m
+	}
+}

v1/ast/external_source_test.go

@@ -0,0 +1,72 @@
+package ast
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+)
+
+type mockExternalSource struct {
+	refs      []Ref
+	rules     []*Rule
+	callCount int32
+}
+
+func newMockExternalSource(refs []Ref, rules []*Rule) *mockExternalSource {
+	return &mockExternalSource{
+		refs:  refs,
+		rules: rules,
+	}
+}
+
+func (m *mockExternalSource) Refs() []Ref {
+	return m.refs
+}
+
+func (m *mockExternalSource) Init(context.Context, Ref) (ExternalRuleIndex, error) {
+	return &mockExternalIndex{rules: m.rules, callCount: &m.callCount}, nil
+}
+
+type mockExternalIndex struct {
+	rules     []*Rule
+	callCount *int32
+}
+
+func (*mockExternalIndex) Opts() *ExternalSourceOptions {
+	return nil
+}
+
+func (m *mockExternalIndex) Lookup(context.Context, ...LookupOption) ([]*Rule, ExternalRuleIndex, error) {
+	atomic.AddInt32(m.callCount, 1)
+	return m.rules, nil, nil
+}
+
+func (m *mockExternalSource) getCallCount() int {
+	return int(atomic.LoadInt32(&m.callCount))
+}
+
+func TestCompilerRuleIndexReturnsNilForExternalSources(t *testing.T) {
+	rule := &Rule{
+		Head: &Head{
+			Reference: MustParseRef("data.external.test.foo"),
+			Value:     BooleanTerm(true),
+		},
+		Body: NewBody(
+			Equality.Expr(VarTerm("x"), IntNumberTerm(1)),
+		),
+	}
+
+	packageRef := MustParseRef("data.external.test")
+	source := newMockExternalSource([]Ref{packageRef}, []*Rule{rule})
+	compiler := NewCompiler()
+	compiler.WithExternalSource(packageRef, source)
+
+	index := compiler.RuleIndex(packageRef)
+	if index != nil {
+		t.Error("Expected RuleIndex to return nil for external source path (delegation to evaluation-time)")
+	}
+
+	if source.getCallCount() != 0 {
+		t.Errorf("Expected GetRules NOT to be called at compile-time, got %d calls", source.getCallCount())
+	}
+}

v1/ast/treenode_dump.go

@@ -0,0 +1,56 @@
+package ast
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+)
+
+// Dump returns a string representation of the tree structure rooted at this node.
+func (n *TreeNode) Dump() string {
+	var sb strings.Builder
+	n.dumpRecursive(&sb, "", "")
+	return sb.String()
+}
+
+func (n *TreeNode) dumpRecursive(sb *strings.Builder, prefix, childPrefix string) {
+	sb.WriteString(prefix)
+	fmt.Fprintf(sb, "%v", n.Key)
+
+	if n.Hide {
+		sb.WriteString(" [hidden]")
+	}
+	if n.External != nil {
+		fmt.Fprintf(sb, " ext:%v", n.External.Ref)
+	}
+	if len(n.Values) > 0 {
+		fmt.Fprintf(sb, " rules:%d", len(n.Values))
+	}
+	sb.WriteString("\n")
+
+	if len(n.Children) == 0 {
+		return
+	}
+
+	keys := make([]Value, 0, len(n.Children))
+	for k := range n.Children {
+		keys = append(keys, k)
+	}
+	sort.Slice(keys, func(i, j int) bool {
+		return Compare(keys[i], keys[j]) < 0
+	})
+
+	for i, key := range keys {
+		child := n.Children[key]
+		isLast := i == len(keys)-1
+		var newPrefix, newChildPrefix string
+		if isLast {
+			newPrefix = childPrefix + "└── "
+			newChildPrefix = childPrefix + "    "
+		} else {
+			newPrefix = childPrefix + "├── "
+			newChildPrefix = childPrefix + "│   "
+		}
+		child.dumpRecursive(sb, newPrefix, newChildPrefix)
+	}
+}

v1/bundle/store.go

@@ -359,8 +359,9 @@ type ActivateOpts struct {
 	TxnCtx                   *storage.Context
 	Compiler                 *ast.Compiler
 	Metrics                  metrics.Metrics
-	Bundles                  map[string]*Bundle     // Optional
-	ExtraModules             map[string]*ast.Module // Optional
+	Bundles                  map[string]*Bundle                               // Optional
+	ExtraModules             map[string]*ast.Module                           // Optional
+	ExternalSources          *util.HasherMap[ast.Ref, ast.ExternalRuleSource] // Optional
 	AuthorizationDecisionRef ast.Ref
 	ParserOptions            ast.ParserOptions
 	Plugin                   string
@@ -535,7 +536,7 @@ func activateBundles(opts *ActivateOpts) error {
 	maps.Copy(remainingAndExtra, remaining)
 	maps.Copy(remainingAndExtra, opts.ExtraModules)
 
-	err = compileModules(opts.Compiler, opts.Metrics, snapshotBundles, remainingAndExtra, opts.legacy, opts.AuthorizationDecisionRef)
+	err = compileModules(opts.Compiler, opts.Metrics, snapshotBundles, remainingAndExtra, opts.legacy, opts.AuthorizationDecisionRef, opts.ExternalSources)
 	if err != nil {
 		return err
 	}
@@ -965,11 +966,19 @@ func writeData(ctx context.Context, store storage.Store, txn storage.Transaction
 	return nil
 }
 
-func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, authorizationDecisionRef ast.Ref) error {
+func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, authorizationDecisionRef ast.Ref, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) error {
 
 	m.Timer(metrics.RegoModuleCompile).Start()
 	defer m.Timer(metrics.RegoModuleCompile).Stop()
 
+	// Apply external sources before compilation
+	if externalSources != nil {
+		externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool {
+			compiler = compiler.WithExternalSource(ref, source)
+			return false
+		})
+	}
+
 	modules := make(map[string]*ast.Module, len(compiler.Modules)+len(extraModules)+len(bundles))
 
 	// preserve any modules already on the compiler
@@ -1000,11 +1009,19 @@ func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[strin
 	return iCompiler.VerifyAuthorizationPolicySchema(compiler, authorizationDecisionRef)
 }
 
-func writeModules(ctx context.Context, store storage.Store, txn storage.Transaction, compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool) error {
+func writeModules(ctx context.Context, store storage.Store, txn storage.Transaction, compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) error {
 
 	m.Timer(metrics.RegoModuleCompile).Start()
 	defer m.Timer(metrics.RegoModuleCompile).Stop()
 
+	// Apply external sources before compilation
+	if externalSources != nil {
+		externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool {
+			compiler = compiler.WithExternalSource(ref, source)
+			return false
+		})
+	}
+
 	modules := map[string]*ast.Module{}
 
 	// preserve any modules already on the compiler

v1/bundle/store_test.go

@@ -3714,7 +3714,7 @@ func testWriteData(t *testing.T, tc testWriteModuleCase, legacy bool) {
 			}
 		}
 
-		err := writeModules(t.Context(), mockStore, txn, compiler, metrics.NoOp(), tc.bundles, tc.extraMods, legacy)
+		err := writeModules(t.Context(), mockStore, txn, compiler, metrics.NoOp(), tc.bundles, tc.extraMods, legacy, nil)
 		if !tc.expectErr && err != nil {
 			t.Fatalf("unepected error: %s", err)
 		} else if tc.expectErr && err == nil {

v1/compile/compile.go

@@ -87,7 +87,8 @@ type Compiler struct {
 	fsys                         fs.FS                      // file system to use when loading paths
 	ns                           string
 	regoVersion                  ast.RegoVersion
-	followSymlinks               bool // optionally follow symlinks in the bundle directory when building the bundle
+	followSymlinks               bool      // optionally follow symlinks in the bundle directory when building the bundle
+	externalRefs                 []ast.Ref // external entrypoints provided dynamically
 }
 
 // New returns a new compiler instance that can be invoked.
@@ -179,6 +180,12 @@ func (c *Compiler) WithEnablePrintStatements(yes bool) *Compiler {
 	return c
 }
 
+// WithExternalRefs sets the external entrypoints that are provided dynamically.
+func (c *Compiler) WithExternalRefs(refs []ast.Ref) *Compiler {
+	c.externalRefs = refs
+	return c
+}
+
 // WithPaths adds input filepaths to read policy and data from.
 func (c *Compiler) WithPaths(p ...string) *Compiler {
 	c.paths = append(c.paths, p...)

v1/hooks/hooks.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/open-policy-agent/opa/v1/bundle"
 	"github.com/open-policy-agent/opa/v1/config"
 	topdown_cache "github.com/open-policy-agent/opa/v1/topdown/cache"
 )
@@ -44,6 +45,13 @@ func New(hs ...Hook) Hooks {
 	return h
 }
 
+func (hs *Hooks) Append(h Hook) {
+	if hs.m == nil {
+		hs.m = make(map[Hook]struct{})
+	}
+	hs.m[h] = struct{}{}
+}
+
 func (hs Hooks) Each(fn func(Hook)) {
 	for h := range hs.m {
 		fn(h)
@@ -82,13 +90,21 @@ type InterQueryValueCacheHook interface {
 	OnInterQueryValueCache(context.Context, topdown_cache.InterQueryValueCache) error
 }
 
+// BundlePreActivateHook is called before a bundle is activated and its policies
+// are compiled. This allows hooks to inspect the bundle manifest (e.g. metadata)
+// and register external rule sources that will be available during compilation.
+type BundlePreActivateHook interface {
+	OnBundlePreActivate(ctx context.Context, bundleName string, manifest bundle.Manifest) error
+}
+
 func (hs Hooks) Validate() error {
 	for h := range hs.m {
 		switch h.(type) {
 		case InterQueryCacheHook,
 			InterQueryValueCacheHook,
 			ConfigHook,
-			ConfigDiscoveryHook: // OK
+			ConfigDiscoveryHook,
+			BundlePreActivateHook: // OK
 		default:
 			return fmt.Errorf("unknown hook type %T", h)
 		}

v1/plugins/bundle/plugin.go

@@ -24,6 +24,7 @@ import (
 	"github.com/open-policy-agent/opa/v1/ast"
 	"github.com/open-policy-agent/opa/v1/bundle"
 	"github.com/open-policy-agent/opa/v1/download"
+	"github.com/open-policy-agent/opa/v1/hooks"
 	"github.com/open-policy-agent/opa/v1/logging"
 	"github.com/open-policy-agent/opa/v1/metrics"
 	"github.com/open-policy-agent/opa/v1/plugins"
@@ -635,15 +636,26 @@ func (p *Plugin) activate(ctx context.Context, name string, b *bundle.Bundle, is
 
 		var activateErr error
 
+		// Call pre-activation hooks so plugins can inspect the bundle manifest
+		// and register external sources before compilation.
+		p.manager.Hooks().Each(func(h hooks.Hook) {
+			if f, ok := h.(hooks.BundlePreActivateHook); ok {
+				if err := f.OnBundlePreActivate(ctx, name, b.Manifest); err != nil {
+					p.log(name).Warn("Pre-activation hook failed: %v", err)
+				}
+			}
+		})
+
 		opts := &bundle.ActivateOpts{
-			Ctx:           ctx,
-			Store:         p.manager.Store,
-			Txn:           txn,
-			TxnCtx:        params.Context,
-			Compiler:      compiler,
-			Metrics:       p.status[name].Metrics,
-			Bundles:       map[string]*bundle.Bundle{name: b},
-			ParserOptions: p.manager.ParserOptions(),
+			Ctx:             ctx,
+			Store:           p.manager.Store,
+			Txn:             txn,
+			TxnCtx:          params.Context,
+			Compiler:        compiler,
+			Metrics:         p.status[name].Metrics,
+			Bundles:         map[string]*bundle.Bundle{name: b},
+			ExternalSources: p.manager.GetExternalSources(),
+			ParserOptions:   p.manager.ParserOptions(),
 		}
 
 		if p.manager.Info != nil {

v1/plugins/plugins.go

@@ -36,6 +36,7 @@ import (
 	"github.com/open-policy-agent/opa/v1/topdown/cache"
 	"github.com/open-policy-agent/opa/v1/topdown/print"
 	"github.com/open-policy-agent/opa/v1/tracing"
+	"github.com/open-policy-agent/opa/v1/util"
 )
 
 // Factory defines the interface OPA uses to instantiate your plugin.
@@ -239,6 +240,8 @@ type Manager struct {
 	extraMiddlewares             []func(http.Handler) http.Handler
 	extraAuthorizerRoutes        []func(string, []any) bool
 	bundleActivatorPlugin        string
+	externalSources              *util.HasherMap[ast.Ref, ast.ExternalRuleSource]
+	externalSourcesMux           sync.RWMutex
 }
 
 type pluginStatusMsg interface {
@@ -459,6 +462,16 @@ func WithHooks(hs hooks.Hooks) func(*Manager) {
 	}
 }
 
+// Hooks returns the hooks configured on the Manager.
+func (m *Manager) Hooks() hooks.Hooks {
+	return m.hooks
+}
+
+// AppendHook allows adding to the hooks configured on the Manager.
+func (m *Manager) AppendHook(h hooks.Hook) {
+	m.hooks.Append(h)
+}
+
 // WithParserOptions sets the parser options to be used by the plugin manager.
 func WithParserOptions(opts ast.ParserOptions) func(*Manager) {
 	return func(m *Manager) {
@@ -629,6 +642,7 @@ func (m *Manager) Init(ctx context.Context) error {
 			EnablePrintStatements: m.enablePrintStatements,
 			ParserOptions:         m.parserOptions,
 			BundleActivatorPlugin: m.bundleActivatorPlugin,
+			ExternalSources:       m.GetExternalSources(),
 		})
 		if err != nil {
 			return err
@@ -828,6 +842,30 @@ func (m *Manager) setWasmResolvers(rs []*wasm.Resolver) {
 	m.wasmResolvers = rs
 }
 
+// RegisterExternalSource registers an external rule source with the manager.
+// The source will be applied to all compilers created by the manager.
+// This should be called from a plugin's constructor or Start() method.
+func (m *Manager) RegisterExternalSource(pkgRef ast.Ref, source ast.ExternalRuleSource) {
+	m.externalSourcesMux.Lock()
+	defer m.externalSourcesMux.Unlock()
+
+	if m.externalSources == nil {
+		m.externalSources = util.NewHasherMap[ast.Ref, ast.ExternalRuleSource](ast.RefEqual)
+	}
+
+	m.externalSources.Put(pkgRef, source)
+
+	m.logger.Debug("Registered external source for package: %s", pkgRef)
+}
+
+// GetExternalSources returns the registered external sources
+func (m *Manager) GetExternalSources() *util.HasherMap[ast.Ref, ast.ExternalRuleSource] {
+	m.externalSourcesMux.RLock()
+	defer m.externalSourcesMux.RUnlock()
+
+	return m.externalSources
+}
+
 // Start starts the manager. Init() should be called once before Start().
 func (m *Manager) Start(ctx context.Context) error {
 	if m == nil {
@@ -857,6 +895,25 @@ func (m *Manager) Start(ctx context.Context) error {
 		}
 	}
 
+	// After starting plugins, check if any external sources were registered
+	// and recompile if necessary to include them in the rule tree
+	externalSources := m.GetExternalSources()
+	if externalSources != nil && externalSources.Len() > 0 {
+		err := storage.Txn(ctx, m.Store, storage.TransactionParams{}, func(txn storage.Transaction) error {
+			compiler, err := loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions(), externalSources)
+			if err != nil {
+				return err
+			}
+			m.setCompiler(compiler)
+			m.logger.Debug("Recompiled policies with %d external source(s) after plugin startup", externalSources.Len())
+			return nil
+		})
+
+		if err != nil {
+			return fmt.Errorf("failed to recompile with external sources: %w", err)
+		}
+	}
+
 	return nil
 }
 
@@ -1038,7 +1095,7 @@ func (m *Manager) onCommit(ctx context.Context, txn storage.Transaction, event s
 	// compiler on the context but the server does not (nor would users
 	// implementing their own policy loading.)
 	if compiler == nil && event.PolicyChanged() {
-		compiler, _ = loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions())
+		compiler, _ = loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions(), m.GetExternalSources())
 	}
 
 	if compiler != nil {
@@ -1076,7 +1133,7 @@ func (m *Manager) onCommit(ctx context.Context, txn storage.Transaction, event s
 	}
 }
 
-func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, enablePrintStatements bool, popts ast.ParserOptions) (*ast.Compiler, error) {
+func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, enablePrintStatements bool, popts ast.ParserOptions, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) (*ast.Compiler, error) {
 	policies, err := store.ListPolicies(ctx, txn)
 	if err != nil {
 		return nil, err
@@ -1102,6 +1159,14 @@ func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage
 		compiler = compiler.WithDefaultRegoVersion(popts.RegoVersion)
 	}
 
+	// Apply external sources BEFORE compilation
+	if externalSources != nil {
+		externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool {
+			compiler = compiler.WithExternalSource(ref, source)
+			return false
+		})
+	}
+
 	compiler.Compile(modules)
 	return compiler, nil
 }

v1/plugins/plugins_test.go

@@ -14,6 +14,7 @@ import (
 
 	internal_tracing "github.com/open-policy-agent/opa/internal/distributedtracing"
 	"github.com/open-policy-agent/opa/internal/storage/mock"
+	"github.com/open-policy-agent/opa/v1/ast"
 	"github.com/open-policy-agent/opa/v1/logging"
 	"github.com/open-policy-agent/opa/v1/logging/test"
 	"github.com/open-policy-agent/opa/v1/plugins/rest"
@@ -614,3 +615,156 @@ func (p prometheusRegisterMock) Unregister(collector prom.Collector) bool {
 	delete(p.Collectors, collector)
 	return true
 }
+
+// mockExternalSource is a simple implementation for testing
+type mockExternalSource struct {
+	refs  []ast.Ref
+	rules []*ast.Rule
+}
+
+func (m *mockExternalSource) Refs() []ast.Ref {
+	return m.refs
+}
+
+func (m *mockExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &mockExternalIndex{rules: m.rules}, nil
+}
+
+type mockExternalIndex struct {
+	rules []*ast.Rule
+}
+
+func (*mockExternalIndex) Opts() *ast.ExternalSourceOptions {
+	return nil
+}
+
+func (m *mockExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	return m.rules, nil, nil
+}
+
+// testExternalSourcePlugin registers an external source during construction
+type testExternalSourcePlugin struct {
+	manager *Manager
+	started bool
+}
+
+func (p *testExternalSourcePlugin) Start(context.Context) error {
+	p.started = true
+	return nil
+}
+
+func (*testExternalSourcePlugin) Stop(context.Context) {}
+
+func (*testExternalSourcePlugin) Reconfigure(context.Context, any) {}
+
+// TestExternalSourceIntegration verifies external source behavior during plugin lifecycle
+func TestExternalSourceIntegration(t *testing.T) {
+	t.Run("sources wired after plugin start", func(t *testing.T) {
+		ctx := context.Background()
+		m, err := New([]byte(`{}`), "test", inmem.New())
+		if err != nil {
+			t.Fatalf("Failed to create manager: %v", err)
+		}
+
+		if err := m.Init(ctx); err != nil {
+			t.Fatalf("Failed to initialize manager: %v", err)
+		}
+
+		module := ast.MustParseModule(`package external.test
+test_rule := true`)
+		pkgRef := ast.MustParseRef("data.external.test")
+
+		source := &mockExternalSource{
+			refs:  []ast.Ref{pkgRef},
+			rules: module.Rules,
+		}
+
+		plugin := &testExternalSourcePlugin{manager: m}
+		m.Register("test_external_source", plugin)
+		m.RegisterExternalSource(pkgRef, source)
+
+		if m.GetExternalSources() == nil || m.GetExternalSources().Len() != 1 {
+			t.Fatalf("Expected 1 external source, got %d", m.GetExternalSources().Len())
+		}
+
+		if err := m.Start(ctx); err != nil {
+			t.Fatalf("Failed to start manager: %v", err)
+		}
+
+		if !plugin.started {
+			t.Fatal("Expected plugin to be started")
+		}
+
+		compiler := m.GetCompiler()
+		if compiler == nil || compiler.RuleTree == nil {
+			t.Fatal("Expected compiler with rule tree after Start()")
+		}
+	})
+
+	t.Run("stop cleans up external source plugins", func(t *testing.T) {
+		ctx := t.Context()
+		m, err := New([]byte(`{}`), "test", inmem.New())
+		if err != nil {
+			t.Fatalf("Failed to create manager: %v", err)
+		}
+
+		if err := m.Init(ctx); err != nil {
+			t.Fatalf("Failed to initialize manager: %v", err)
+		}
+
+		module := ast.MustParseModule(`package external.test
+test_rule := true`)
+		pkgRef := ast.MustParseRef("data.external.test")
+
+		source := &mockExternalSource{
+			refs:  []ast.Ref{pkgRef},
+			rules: module.Rules,
+		}
+
+		plugin := &testExternalSourcePlugin{manager: m}
+		m.Register("test_external_source", plugin)
+		m.RegisterExternalSource(pkgRef, source)
+
+		if err := m.Start(ctx); err != nil {
+			t.Fatalf("Failed to start manager: %v", err)
+		}
+
+		if !plugin.started {
+			t.Fatal("Expected plugin to be started")
+		}
+
+		m.Stop(ctx)
+
+		if m.GetExternalSources() == nil || m.GetExternalSources().Len() != 1 {
+			t.Fatalf("Expected external sources to still be registered after stop, got %d", m.GetExternalSources().Len())
+		}
+	})
+
+	t.Run("no recompilation when no sources registered", func(t *testing.T) {
+		ctx := context.Background()
+		m, err := New([]byte(`{}`), "test", inmem.New())
+		if err != nil {
+			t.Fatalf("Failed to create manager: %v", err)
+		}
+
+		if err := m.Init(ctx); err != nil {
+			t.Fatalf("Failed to initialize manager: %v", err)
+		}
+
+		compilerBeforeStart := m.GetCompiler()
+		if compilerBeforeStart == nil {
+			t.Fatal("Expected compiler to be initialized after Init()")
+		}
+
+		plugin := &testPlugin{m: m}
+		m.Register("test_plugin", plugin)
+
+		if err := m.Start(ctx); err != nil {
+			t.Fatalf("Failed to start manager: %v", err)
+		}
+
+		if m.GetCompiler() != compilerBeforeStart {
+			t.Fatal("Expected compiler to remain the same when no external sources registered")
+		}
+	})
+}

v1/rego/rego.go

@@ -673,6 +673,7 @@ type Rego struct {
 	strictBuiltinErrors         bool
 	builtinErrorList            *[]topdown.Error
 	resolvers                   []refResolver
+	externalSources             []ast.ExternalRuleSource
 	schemaSet                   *ast.SchemaSet
 	target                      string // target type (wasm, rego, etc.)
 	opa                         opa.EvalEngine
@@ -1287,6 +1288,15 @@ func Resolver(ref ast.Ref, r resolver.Resolver) func(r *Rego) {
 	}
 }
 
+// ExternalSource adds an external rule source that provides rules dynamically.
+// The source declares which package refs it handles via its Refs() method.
+// A single source can provide rules for multiple packages.
+func ExternalSource(source ast.ExternalRuleSource) func(r *Rego) {
+	return func(rego *Rego) {
+		rego.externalSources = append(rego.externalSources, source)
+	}
+}
+
 // Schemas sets the schemaSet
 func Schemas(x *ast.SchemaSet) func(r *Rego) {
 	return func(r *Rego) {
@@ -2142,6 +2152,18 @@ func parserOptionsFromRegoVersionImport(imports []*ast.Import, popts ast.ParserO
 }
 
 func (r *Rego) compileModules(ctx context.Context, txn storage.Transaction, m metrics.Metrics) error {
+	if len(r.externalSources) > 0 && r.target != "" && r.target != targetRego {
+		return fmt.Errorf("external rule sources are not supported with target %q: only the default (rego) target is supported", r.target)
+	}
+
+	// Apply external sources to the compiler before compilation
+	for i := range r.externalSources {
+		source := r.externalSources[i]
+		for _, ref := range source.Refs() {
+			r.compiler.WithExternalSource(ref, source)
+		}
+	}
+
 	// Only compile again if there are new modules.
 	if len(r.bundles) > 0 || len(r.parsedModules) > 0 {
 

v1/rego/rego_external_source_test.go

@@ -0,0 +1,456 @@
+package rego
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/open-policy-agent/opa/v1/ast"
+)
+
+type mockExternalSource struct {
+	refRules    map[string][]*ast.Rule
+	visibleRefs []ast.Ref
+}
+
+func newMockExternalSource(refs []ast.Ref, rules []*ast.Rule) *mockExternalSource {
+	refRules := make(map[string][]*ast.Rule)
+	for _, ref := range refs {
+		refRules[ref.String()] = rules
+	}
+	return &mockExternalSource{
+		refRules: refRules,
+	}
+}
+
+func (m *mockExternalSource) Refs() []ast.Ref {
+	refs := make([]ast.Ref, 0, len(m.refRules))
+	for refStr := range m.refRules {
+		refs = append(refs, ast.MustParseRef(refStr))
+	}
+	return refs
+}
+
+func (m *mockExternalSource) Init(_ context.Context, ref ast.Ref) (ast.ExternalRuleIndex, error) {
+	rules, ok := m.refRules[ref.String()]
+	if !ok {
+		return nil, nil
+	}
+	return &mockExternalIndex{rules: rules, visibleRefs: m.visibleRefs}, nil
+}
+
+type mockExternalIndex struct {
+	rules       []*ast.Rule
+	visibleRefs []ast.Ref
+}
+
+func (m *mockExternalIndex) Opts() *ast.ExternalSourceOptions {
+	if m.visibleRefs == nil {
+		return nil
+	}
+	return &ast.ExternalSourceOptions{VisibleRefs: m.visibleRefs}
+}
+
+func (m *mockExternalIndex) Lookup(ctx context.Context, _ ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	return m.rules, nil, nil
+}
+
+func evalWithExternalSource(t *testing.T, ctx context.Context, query, module string, source *mockExternalSource, input map[string]any) ResultSet {
+	t.Helper()
+	opts := []func(*Rego){
+		Query(query),
+		Module("test.rego", module),
+		ExternalSource(source),
+	}
+	if input != nil {
+		opts = append(opts, Input(input))
+	}
+	r := New(opts...)
+	rs, err := r.Eval(ctx)
+	if err != nil {
+		t.Fatalf("Eval failed: %v", err)
+	}
+	return rs
+}
+
+func partialEvalWithExternalSource(t *testing.T, ctx context.Context, query, module string, source *mockExternalSource) *PartialQueries {
+	t.Helper()
+	r := New(
+		Query(query),
+		Module("test.rego", module),
+		ExternalSource(source),
+	)
+	pq, err := r.Partial(ctx)
+	if err != nil {
+		t.Fatalf("Partial failed: %v", err)
+	}
+	return pq
+}
+
+func assertBoolResult(t *testing.T, rs ResultSet, expected bool, msg string) {
+	t.Helper()
+	if len(rs) != 1 {
+		t.Fatalf("Expected 1 result, got %d", len(rs))
+	}
+	result, ok := rs[0].Expressions[0].Value.(bool)
+	if !ok {
+		t.Fatalf("Expected boolean result, got %T", rs[0].Expressions[0].Value)
+	}
+	if result != expected {
+		t.Errorf("%s: expected %v, got %v", msg, expected, result)
+	}
+}
+
+func assertNoResults(t *testing.T, rs ResultSet, msg string) {
+	t.Helper()
+	if len(rs) != 0 {
+		t.Errorf("%s: expected 0 results, got %d", msg, len(rs))
+	}
+}
+
+func assertPartialQuery(t *testing.T, pq *PartialQueries, expectedQueries string) {
+	t.Helper()
+	expected := ast.MustParseBody(expectedQueries)
+	if len(pq.Queries) != 1 {
+		t.Fatalf("Expected 1 query, got %d", len(pq.Queries))
+	}
+	if !pq.Queries[0].Equal(expected) {
+		t.Errorf("Expected PE result:\n%v\n\nGot:\n%v", expected, pq.Queries[0])
+	}
+}
+
+func TestExternalSourceDecisionMaking(t *testing.T) {
+	ctx := t.Context()
+
+	externalModule := ast.MustParseModule(`package external.authz
+allow if input.role == "admin"`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+	source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules)
+
+	staticModule := `package authz
+default allow := false
+allow if data.external.authz.allow`
+
+	t.Run("admin role allowed", func(t *testing.T) {
+		rs := evalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source, map[string]any{"role": "admin"})
+		assertBoolResult(t, rs, true, "Expected allow=true for admin role")
+	})
+
+	t.Run("non-admin role not allowed", func(t *testing.T) {
+		rs := evalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source, map[string]any{"role": "user"})
+		assertBoolResult(t, rs, false, "Expected allow=false for user role")
+	})
+}
+
+func TestExternalSourcePartialEval(t *testing.T) {
+	ctx := t.Context()
+
+	externalModule := ast.MustParseModule(`package external.authz
+allow if input.role == "admin"`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+	source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules)
+
+	staticModule := `package authz
+default allow := false
+allow if data.external.authz.allow`
+
+	t.Run("partial eval into external rule", func(t *testing.T) {
+		pq := partialEvalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source)
+		assertPartialQuery(t, pq, `input.role = "admin"`)
+	})
+}
+
+func TestExternalSourceCallBackIntoStaticRego(t *testing.T) {
+	ctx := t.Context()
+
+	externalModule := ast.MustParseModule(`package external.authz
+allow if data.static.authz.foo == "bar"`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+	source := &mockExternalSource{
+		refRules: map[string][]*ast.Rule{
+			packageRef.String(): externalModule.Rules,
+		},
+		visibleRefs: []ast.Ref{ast.MustParseRef("data")},
+	}
+
+	staticModule := `package static.authz
+default allow := false
+allow if data.external.authz.allow
+
+foo := "bar"`
+
+	rs := evalWithExternalSource(t, ctx, "data.static.authz.allow", staticModule, source, nil)
+	assertBoolResult(t, rs, true, "Expected allow=true")
+}
+
+func TestExternalSourceCallBackIntoStaticRegoWithRecursion(t *testing.T) {
+	externalModule := ast.MustParseModule(`package external.authz
+allow if data.static.authz.allow`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+
+	staticModule := `package static.authz
+default allow := false
+allow if data.external.authz.allow`
+
+	t.Run("isolated by default prevents recursion", func(t *testing.T) {
+		ctx := t.Context()
+		source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules)
+		rs := evalWithExternalSource(t, ctx, "data.static.authz.allow", staticModule, source, nil)
+		assertBoolResult(t, rs, false, "Expected allow=false (isolated external source cannot access static policy)")
+	})
+
+	t.Run("visible refs allows recursion and hits deadline", func(t *testing.T) {
+		ctx := t.Context()
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				packageRef.String(): externalModule.Rules,
+			},
+			visibleRefs: []ast.Ref{ast.MustParseRef("data")},
+		}
+
+		r := New(
+			Query("data.static.authz.allow"),
+			Module("authz.rego", staticModule),
+			ExternalSource(source),
+		)
+
+		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		t.Cleanup(cancel)
+		_, err := r.Eval(ctx)
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Errorf("expected deadline-exceeded, got err: %v", err)
+		}
+	})
+}
+
+func TestExternalSourceCrossRefCalls(t *testing.T) {
+	ctx := t.Context()
+
+	fooModule := ast.MustParseModule(`package external.foo
+result if data.external.bar.value == 42`)
+
+	barModule := ast.MustParseModule(`package external.bar
+value := 42`)
+
+	staticModule := `package main
+result := data.external.foo.result`
+
+	t.Run("isolated prevents cross-ref calls", func(t *testing.T) {
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				"data.external.foo": fooModule.Rules,
+				"data.external.bar": barModule.Rules,
+			},
+			visibleRefs: []ast.Ref{}, // explicitly isolated
+		}
+
+		rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, nil)
+		assertNoResults(t, rs, "Expected 0 results when isolated (external.foo cannot access external.bar)")
+	})
+
+	t.Run("visible refs allows cross-ref calls", func(t *testing.T) {
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				"data.external.foo": fooModule.Rules,
+				"data.external.bar": barModule.Rules,
+			},
+			visibleRefs: []ast.Ref{ast.MustParseRef("data")},
+		}
+
+		rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, nil)
+		assertBoolResult(t, rs, true, "Expected result=true when all refs visible")
+	})
+}
+
+func TestExternalSourceNestedPackage(t *testing.T) {
+	ctx := t.Context()
+
+	externalModule := ast.MustParseModule(`package external.project.authz
+
+allowed if other_rule
+other_rule if input.foo == "bar"`)
+
+	parentRef := ast.MustParseRef("data.external")
+	source := newMockExternalSource([]ast.Ref{parentRef}, externalModule.Rules)
+
+	staticModule := `package main
+result := data.external.project.authz.allowed`
+
+	t.Run("eval allowed when foo is bar", func(t *testing.T) {
+		rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, map[string]any{"foo": "bar"})
+		assertBoolResult(t, rs, true, "Expected allowed=true when foo is bar")
+	})
+
+	t.Run("eval denied when foo is not bar", func(t *testing.T) {
+		rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, map[string]any{"foo": "baz"})
+		assertNoResults(t, rs, "Expected no results when foo is not bar")
+	})
+
+	t.Run("partial eval", func(t *testing.T) {
+		pq := partialEvalWithExternalSource(t, ctx, "data.main.result", staticModule, source)
+		assertPartialQuery(t, pq, `input.foo = "bar"`)
+	})
+
+	t.Run("call external function with argument fails", func(t *testing.T) {
+		externalModuleWithFunc := ast.MustParseModule(`package external.authz
+
+foo(x) if x == "bar"`)
+
+		funcSource := newMockExternalSource([]ast.Ref{ast.MustParseRef("data.external.authz")}, externalModuleWithFunc.Rules)
+
+		staticModuleWithFuncCall := `package main
+allow if data.external.authz.foo("bar")`
+
+		r := New(
+			Query("data.main.allow"),
+			Module("test.rego", staticModuleWithFuncCall),
+			ExternalSource(funcSource),
+		)
+		_, err := r.Eval(ctx)
+		if err == nil {
+			t.Fatal("Expected error when calling external function with argument, but got none")
+		}
+		var errs ast.Errors
+		if !errors.As(err, &errs) {
+			t.Fatalf("Expected ast.Errors, got: %T: %v", err, err)
+		}
+		if !slices.ContainsFunc(errs, func(e *ast.Error) bool { return e.Code == ast.TypeErr }) {
+			t.Errorf("Expected type error for undefined function, got: %v", err)
+		}
+	})
+
+	t.Run("call into external rule that uses function internally", func(t *testing.T) {
+		externalModuleWithInternalFunc := ast.MustParseModule(`package external.authz
+
+allowed if foo("bar")
+foo(x) if x == "bar"`)
+
+		funcSource := newMockExternalSource([]ast.Ref{ast.MustParseRef("data.external.authz")}, externalModuleWithInternalFunc.Rules)
+
+		staticModuleWithRuleCall := `package main
+allow if data.external.authz.allowed`
+
+		rs := evalWithExternalSource(t, ctx, "data.main.allow", staticModuleWithRuleCall, funcSource, nil)
+		assertBoolResult(t, rs, true, "Expected allow=true when external rule internally uses function")
+	})
+}
+
+func TestExternalSourcePartialVisibility(t *testing.T) {
+	ctx := t.Context()
+
+	externalModule := ast.MustParseModule(`package external.authz
+allow if {
+	data.pkg_a.check
+	data.pkg_b.check
+}`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+
+	modA := `package pkg_a
+check if input.role == "admin"`
+
+	modB := `package pkg_b
+check := true`
+
+	t.Run("only pkg_a visible", func(t *testing.T) {
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				packageRef.String(): externalModule.Rules,
+			},
+			visibleRefs: []ast.Ref{ast.MustParseRef("data.pkg_a")},
+		}
+
+		r := New(
+			Query("data.external.authz.allow"),
+			Module("a.rego", modA),
+			Module("b.rego", modB),
+			ExternalSource(source),
+			Input(map[string]any{"role": "admin"}),
+		)
+		rs, err := r.Eval(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertNoResults(t, rs, "Expected 0 results (pkg_b not visible)")
+	})
+
+	t.Run("both visible", func(t *testing.T) {
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				packageRef.String(): externalModule.Rules,
+			},
+			visibleRefs: []ast.Ref{
+				ast.MustParseRef("data.pkg_a"),
+				ast.MustParseRef("data.pkg_b"),
+			},
+		}
+
+		r := New(
+			Query("data.external.authz.allow"),
+			Module("a.rego", modA),
+			Module("b.rego", modB),
+			ExternalSource(source),
+			Input(map[string]any{"role": "admin"}),
+		)
+		rs, err := r.Eval(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertBoolResult(t, rs, true, "Expected allow=true when both packages visible")
+	})
+
+	t.Run("all of data visible", func(t *testing.T) {
+		source := &mockExternalSource{
+			refRules: map[string][]*ast.Rule{
+				packageRef.String(): externalModule.Rules,
+			},
+			visibleRefs: []ast.Ref{ast.MustParseRef("data")},
+		}
+
+		r := New(
+			Query("data.external.authz.allow"),
+			Module("a.rego", modA),
+			Module("b.rego", modB),
+			ExternalSource(source),
+			Input(map[string]any{"role": "admin"}),
+		)
+		rs, err := r.Eval(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertBoolResult(t, rs, true, "Expected allow=true when all of data is visible")
+	})
+}
+
+func TestExternalSourceUnsupportedTarget(t *testing.T) {
+	externalModule := ast.MustParseModule(`package external.authz
+allow if input.role == "admin"`)
+
+	packageRef := ast.MustParseRef("data.external.authz")
+	source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules)
+
+	for _, target := range []string{"wasm", "plan", "custom-plugin"} {
+		t.Run(target, func(t *testing.T) {
+			r := New(
+				Query("data.external.authz.allow"),
+				Module("test.rego", `package test`),
+				ExternalSource(source),
+				Target(target),
+			)
+			_, err := r.Eval(t.Context())
+			if err == nil {
+				t.Fatal("Expected error for non-rego target with external sources")
+			}
+			if !strings.Contains(err.Error(), "external rule sources are not supported") {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+		})
+	}
+}

v1/topdown/eval.go

@@ -92,8 +92,9 @@ type eval struct {
 	input                       *ast.Term
 	data                        *ast.Term
 	external                    *resolverTrie
+	externalTreeStack           *externalTreeStack
 	targetStack                 *refStack
-	traceLastLocation           *ast.Location // Last location of a trace event.
+	traceLastLocation           *ast.Location
 	instr                       *Instrumentation
 	builtins                    map[string]*Builtin
 	builtinCache                builtins.Cache
@@ -267,12 +268,15 @@ func (e *eval) unknown(x any, b *bindings) bool {
 		x = ast.NewTerm(v)
 	}
 
-	return saveRequired(e.compiler, e.inliningControl, true, e.saveSet, b, x, false)
+	return saveRequired(e.compiler.RuleTree, e.externalTreeStack, e.inliningControl, true, e.saveSet, b, x, false)
 }
 
 // exactly like `unknown` above` but without the cost of `any` boxing when arg is known to be a ref
 func (e *eval) unknownRef(ref ast.Ref, b *bindings) bool {
-	return e.partial() && saveRequired(e.compiler, e.inliningControl, true, e.saveSet, b, ast.NewTerm(ref), false)
+	if !e.partial() {
+		return false
+	}
+	return saveRequired(e.compiler.RuleTree, e.externalTreeStack, e.inliningControl, true, e.saveSet, b, ast.NewTerm(ref), false)
 }
 
 func (e *eval) traceEnter(x ast.Node) {
@@ -755,26 +759,33 @@ func (e *eval) evalWith(iter evalIterator) error {
 		}
 	}
 
-	oldInput, oldData := e.evalWithPush(input, data, functionMocks, targets, disable)
+	oldInput, oldData, pushedFrame := e.evalWithPush(input, data, functionMocks, targets, disable)
 
 	err = e.evalStep(func(e *eval) error {
-		e.evalWithPop(oldInput, oldData)
+		e.evalWithPop(oldInput, oldData, pushedFrame)
 		err := e.next(iter)
-		oldInput, oldData = e.evalWithPush(input, data, functionMocks, targets, disable)
+		oldInput, oldData, pushedFrame = e.evalWithPush(input, data, functionMocks, targets, disable)
 		return err
 	})
 
-	e.evalWithPop(oldInput, oldData)
+	e.evalWithPop(oldInput, oldData, pushedFrame)
 
 	return err
 }
 
-func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term, targets, disable []ast.Ref) (*ast.Term, *ast.Term) {
+func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term, targets, disable []ast.Ref) (*ast.Term, *ast.Term, bool) {
 	var oldInput *ast.Term
+	var pushedFrame bool
 
 	if input != nil {
 		oldInput = e.input
 		e.input = input
+
+		// When input changes, push a new frame for external tree caching
+		if e.externalTreeStack != nil {
+			e.externalTreeStack.PushFrame()
+			pushedFrame = true
+		}
 	}
 
 	var oldData *ast.Term
@@ -804,16 +815,22 @@ func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term,
 
 	e.functionMocks.PutPairs(functionMocks)
 
-	return oldInput, oldData
+	return oldInput, oldData, pushedFrame
 }
 
-func (e *eval) evalWithPop(input, data *ast.Term) {
+func (e *eval) evalWithPop(input, data *ast.Term, popFrame bool) {
 	// NOTE(ae) no nil checks here as we assume evalWithPush always called first
 	e.inliningControl.PopDisable()
 	e.targetStack.Pop()
 	e.virtualCache.Pop()
 	e.comprehensionCache.Pop()
 	e.functionMocks.PopPairs()
+
+	// When input is restored, pop the external tree frame
+	if popFrame {
+		e.externalTreeStack.PopFrame()
+	}
+
 	e.data = data
 	e.input = input
 }
@@ -980,15 +997,19 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 
 		var ir *ast.IndexResult
 		var err error
+		index := e.ruleIndex(ref)
 		if e.partial() {
-			ir, err = e.getRules(ref, nil)
+			ir, err = e.getRules(ref, nil, index)
 		} else {
-			ir, err = e.getRules(ref, terms[1:])
+			ir, err = e.getRules(ref, terms[1:], index)
 		}
 		defer ast.IndexResultPool.Put(ir)
 		if err != nil {
 			return err
 		}
+		if ir == nil {
+			return nil
+		}
 
 		eval := evalFuncPool.Get()
 		defer evalFuncPool.Put(eval)
@@ -1724,11 +1745,10 @@ func (e *eval) saveInlinedNegatedExprs(exprs []*ast.Expr, iter unifyIterator) er
 	return err
 }
 
-func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error) {
+func (e *eval) getRules(ref ast.Ref, args []*ast.Term, index ast.RuleIndex) (*ast.IndexResult, error) {
 	e.instr.startTimer(evalOpRuleIndex)
 	defer e.instr.stopTimer(evalOpRuleIndex)
 
-	index := e.ruleIndex(ref)
 	if index == nil {
 		return nil, nil
 	}
@@ -1742,6 +1762,7 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error)
 
 	var result *ast.IndexResult
 	var err error
+
 	resolver.e = e
 	if e.indexing {
 		resolver.args = args
@@ -1779,11 +1800,6 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error)
 	return result, err
 }
 
-// ruleIndex performs a lookup for a RuleIndex in the compiler's RuleTree.
-func (e *eval) ruleIndex(ref ast.Ref) ast.RuleIndex {
-	return e.compiler.RuleIndex(ref)
-}
-
 func (e *eval) Resolve(ref ast.Ref) (ast.Value, error) {
 	return (&evalResolver{e: e}).Resolve(ref)
 }
@@ -1996,7 +2012,7 @@ func (e *eval) getDeclArgsLen(x *ast.Expr) (int, error) {
 		return bi.Decl.Arity(), nil
 	}
 
-	ir, err := e.getRules(operator, nil)
+	ir, err := e.getRules(operator, nil, e.ruleIndex(operator))
 	defer ast.IndexResultPool.Put(ir)
 	if err != nil {
 		return -1, err
@@ -2536,10 +2552,48 @@ func (e evalTree) next(iter unifyIterator, plugged *ast.Term) error {
 	cpy.plugged[e.pos] = plugged
 	cpy.pos++
 
+	// Track whether we pushed an external tree that needs cleanup
+	pushedExternalTree := false
 	if !e.e.targetStack.Prefixed(cpy.plugged[:cpy.pos]) {
 		if e.node != nil {
 			node = e.node.Child(plugged.Value)
-			if node != nil && len(node.Values) > 0 {
+
+			// Handle external sources transparently
+			if node != nil && node.External != nil {
+				externalRef := node.External.Ref
+				externalIndex := node.External.Index
+
+				// Initialize externalTreeStack if needed
+				if e.e.externalTreeStack == nil {
+					e.e.externalTreeStack = newExternalTreeStack(e.e)
+				}
+
+				// Check cache first
+				cachedNode, _, found := e.e.externalTreeStack.findCached(externalRef)
+				if found {
+					node = cachedNode
+				} else {
+					// Call Tree() and cache the result
+					e.e.instr.startTimer(evalOpExternalRuleSource)
+					tree, updatedIndex, err := node.External.Tree(e.e.ctx, e.e.compiler.RuleTree, externalRef, e.e.input, e.e.metrics, e.e.requestMetadata, e.e.responseMetadata)
+					e.e.instr.stopTimer(evalOpExternalRuleSource)
+					if err != nil {
+						return err
+					}
+					if tree != nil {
+						if updatedIndex != nil {
+							externalIndex = updatedIndex
+						}
+						e.e.externalTreeStack.Push(externalRef, tree, externalIndex, e.e.input)
+						node = tree
+						pushedExternalTree = true
+					}
+				}
+			}
+
+			hasRules := node != nil && len(node.Values) > 0
+
+			if hasRules {
 				r := evalVirtual{
 					e:         e.e,
 					ref:       e.ref,
@@ -2550,13 +2604,21 @@ func (e evalTree) next(iter unifyIterator, plugged *ast.Term) error {
 					rbindings: e.rbindings,
 				}
 				r.plugged[e.pos] = plugged
-				return r.eval(iter)
+				err := r.eval(iter)
+				if pushedExternalTree {
+					e.e.externalTreeStack.Pop()
+				}
+				return err
 			}
 		}
 	}
 
 	cpy.node = node
-	return cpy.eval(iter)
+	err := cpy.eval(iter)
+	if pushedExternalTree {
+		e.e.externalTreeStack.Pop()
+	}
+	return err
 }
 
 // enumerateNext is a helper to avoid closure allocation in enumerate loops.
@@ -2732,12 +2794,16 @@ type evalVirtual struct {
 
 func (e evalVirtual) eval(iter unifyIterator) error {
 
-	ir, err := e.e.getRules(e.plugged[:e.pos+1], nil)
+	ir, err := e.e.getRules(e.plugged[:e.pos+1], nil, e.e.ruleIndex(e.plugged[:e.pos+1]))
 	defer ast.IndexResultPool.Put(ir)
 	if err != nil {
 		return err
 	}
 
+	if ir == nil {
+		return nil
+	}
+
 	// Partial evaluation of ordered rules is not supported currently. Save the
 	// expression and continue. This could be revisited in the future.
 	if len(ir.Else) > 0 && e.e.unknownRef(e.ref, e.bindings) {
@@ -4490,3 +4556,180 @@ func (e *eval) updateSavedMocks(withs []*ast.With) []*ast.With {
 	}
 	return ret
 }
+
+// simpleTreeNode provides minimal tree structure for navigation
+type simpleTreeNode struct {
+	tree     *ast.TreeNode
+	children map[ast.Value]*simpleTreeNode
+}
+
+func newSimpleTreeNode() *simpleTreeNode {
+	return &simpleTreeNode{
+		children: make(map[ast.Value]*simpleTreeNode),
+	}
+}
+
+// externalTreeStack caches external rule trees and tracks frames for input changes.
+// It maintains both a flat cache for lookups and a tree structure for navigation.
+type externalTreeStack struct {
+	eval    *eval
+	entries []externalTreeEntry // flat list of cached entries
+	frames  []int               // frame markers for input changes (indices into entries)
+	root    *simpleTreeNode     // tree structure for navigation
+}
+
+type externalTreeEntry struct {
+	ref   ast.Ref
+	input *ast.Term
+	tree  *ast.TreeNode
+	index ast.ExternalRuleIndex
+}
+
+func newExternalTreeStack(e *eval) *externalTreeStack {
+	return &externalTreeStack{
+		eval:    e,
+		entries: make([]externalTreeEntry, 0, 4),
+		frames:  make([]int, 0, 4),
+	}
+}
+
+// findCached checks if we already have a cached tree for this ref.
+// Frame tracking ensures any cached entry has the correct input.
+func (s *externalTreeStack) findCached(ref ast.Ref) (*ast.TreeNode, ast.ExternalRuleIndex, bool) {
+	// Determine search boundary: only search within current frame if one exists
+	startIdx := 0
+	if len(s.frames) > 0 {
+		startIdx = s.frames[len(s.frames)-1]
+	}
+
+	// Search from most recent to the frame boundary
+	for i := len(s.entries) - 1; i >= startIdx; i-- {
+		entry := &s.entries[i]
+		if entry.ref.Equal(ref) {
+			return entry.tree, entry.index, true
+		}
+	}
+	return nil, nil, false
+}
+
+func (s *externalTreeStack) Push(ref ast.Ref, tree *ast.TreeNode, index ast.ExternalRuleIndex, input *ast.Term) {
+	// Add entry to cache (we never have duplicates in the same frame)
+	s.entries = append(s.entries, externalTreeEntry{
+		ref:   ref,
+		input: input,
+		tree:  tree,
+		index: index,
+	})
+
+	// Update root tree structure
+	if s.root == nil {
+		s.root = newSimpleTreeNode()
+	}
+	node := s.root
+	for _, term := range ref {
+		key := term.Value
+		if node.children[key] == nil {
+			node.children[key] = newSimpleTreeNode()
+		}
+		node = node.children[key]
+	}
+	node.tree = tree
+}
+
+// Pop removes the most recent entry from the stack and closes its index.
+// If a frame is active and entries are at or below the frame boundary,
+// Pop is a no-op — PopFrame already cleaned up (or will clean up) those entries.
+func (s *externalTreeStack) Pop() {
+	if len(s.entries) == 0 {
+		return
+	}
+
+	// Don't pop at or below the current frame boundary.
+	if len(s.frames) > 0 && len(s.entries) <= s.frames[len(s.frames)-1] {
+		return
+	}
+
+	// Close the external index if it supports closing
+	lastEntry := &s.entries[len(s.entries)-1]
+	if closer, ok := lastEntry.index.(ast.ExternalRuleIndexCloser); ok {
+		_ = closer.Close()
+	}
+
+	// Remove the most recent entry
+	s.entries = s.entries[:len(s.entries)-1]
+
+	// Rebuild tree from remaining entries
+	s.root = newSimpleTreeNode()
+	for i := range s.entries {
+		entry := &s.entries[i]
+		node := s.root
+		for _, term := range entry.ref {
+			key := term.Value
+			if node.children[key] == nil {
+				node.children[key] = newSimpleTreeNode()
+			}
+			node = node.children[key]
+		}
+		node.tree = entry.tree
+	}
+}
+
+// PushFrame marks the current stack position for input changes
+func (s *externalTreeStack) PushFrame() {
+	s.frames = append(s.frames, len(s.entries))
+}
+
+// PopFrame restores the cache to the marked position
+func (s *externalTreeStack) PopFrame() {
+	if len(s.frames) == 0 {
+		return
+	}
+
+	// Get the frame marker and truncate entries
+	targetSize := s.frames[len(s.frames)-1]
+	s.frames = s.frames[:len(s.frames)-1]
+
+	// Close indices of entries being removed
+	for i := len(s.entries) - 1; i >= targetSize; i-- {
+		if closer, ok := s.entries[i].index.(ast.ExternalRuleIndexCloser); ok {
+			_ = closer.Close()
+		}
+	}
+
+	// Rebuild tree from remaining entries
+	s.root = newSimpleTreeNode()
+	for i := range targetSize {
+		entry := &s.entries[i]
+		node := s.root
+		for _, term := range entry.ref {
+			key := term.Value
+			if node.children[key] == nil {
+				node.children[key] = newSimpleTreeNode()
+			}
+			node = node.children[key]
+		}
+		node.tree = entry.tree
+	}
+
+	s.entries = s.entries[:targetSize]
+}
+
+// ruleIndex performs a shadowed lookup for a RuleIndex, checking external trees first.
+// It searches through the pushStack (most recent to oldest), navigating the tree structure
+// to find matching rules, then falls back to the compiler's static RuleTree.
+func (e *eval) ruleIndex(ref ast.Ref) ast.RuleIndex {
+	if e.externalTreeStack != nil && len(e.externalTreeStack.entries) > 0 {
+		// Search from most recent to oldest
+		for i := len(e.externalTreeStack.entries) - 1; i >= 0; i-- {
+			entry := &e.externalTreeStack.entries[i]
+			if ref.HasPrefix(entry.ref) {
+				// Look for the relative ref in the cached tree
+				relativeRef := ref[len(entry.ref):]
+				if found := entry.tree.Find(relativeRef); found != nil {
+					return found.Index
+				}
+			}
+		}
+	}
+	return e.compiler.RuleIndex(ref)
+}

v1/topdown/external_source_test.go

@@ -0,0 +1,429 @@
+package topdown
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"sync/atomic"
+	"testing"
+
+	"github.com/open-policy-agent/opa/v1/ast"
+	"github.com/open-policy-agent/opa/v1/metrics"
+	"github.com/open-policy-agent/opa/v1/storage/inmem"
+)
+
+type countingExternalSource struct {
+	refs      []ast.Ref
+	rules     []*ast.Rule
+	callCount int32
+}
+
+func (m *countingExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &countingExternalIndex{rules: m.rules, callCount: &m.callCount}, nil
+}
+
+func (m *countingExternalSource) Refs() []ast.Ref {
+	return m.refs
+}
+
+type countingExternalIndex struct {
+	rules     []*ast.Rule
+	callCount *int32
+}
+
+func (*countingExternalIndex) Opts() *ast.ExternalSourceOptions {
+	return nil
+}
+
+func (m *countingExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	atomic.AddInt32(m.callCount, 1)
+	return m.rules, nil, nil
+}
+
+func (m *countingExternalSource) getCallCount() int {
+	return int(atomic.LoadInt32(&m.callCount))
+}
+
+func setupCompiler(t *testing.T, packageRef ast.Ref, source ast.ExternalRuleSource, staticModule *ast.Module) *ast.Compiler {
+	t.Helper()
+	compiler := ast.NewCompiler()
+	compiler.WithExternalSource(packageRef, source)
+	modules := map[string]*ast.Module{}
+	if staticModule != nil {
+		modules["main.rego"] = staticModule
+	}
+	compiler.Compile(modules)
+	if compiler.Failed() {
+		t.Fatalf("Compiler failed: %v", compiler.Errors)
+	}
+	return compiler
+}
+
+func runQuery(t *testing.T, compiler *ast.Compiler, queryStr string, input *ast.Term) QueryResultSet {
+	t.Helper()
+	store := inmem.New()
+	ctx := t.Context()
+	txn, err := store.NewTransaction(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Abort(ctx, txn)
+	m := metrics.New()
+	instr := NewInstrumentation(m)
+
+	query := ast.MustParseBody(queryStr)
+	q := NewQuery(query).
+		WithCompiler(compiler).
+		WithStore(store).
+		WithTransaction(txn).
+		WithInput(input).
+		WithInstrumentation(instr)
+
+	qrs, err := q.Run(ctx)
+	if err != nil {
+		t.Fatalf("Query failed: %v", err)
+	}
+	t.Logf("metrics: %v", m.All())
+	return qrs
+}
+
+func TestExternalSourceE2EWithInputOverride(t *testing.T) {
+	t.Parallel()
+
+	externalModule := ast.MustParseModule(`package authz
+allowed if input.user == "alice"`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	staticModule := ast.MustParseModule(`package main
+check if {
+	data.authz.allowed
+	data.authz.allowed with input as {"user": "bob"}
+}`)
+
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	input := ast.MustParseTerm(`{"user": "alice"}`)
+	qrs := runQuery(t, compiler, "data.main.check", input)
+
+	if len(qrs) != 0 {
+		t.Errorf("Expected 0 results (second check with bob should fail), got %d", len(qrs))
+	}
+
+	if callCount := source.getCallCount(); callCount != 2 {
+		t.Errorf("Expected external source to be called twice (once per input), got %d calls", callCount)
+	}
+}
+
+func TestExternalSourceE2EWithMultipleRulesFromSamePackage(t *testing.T) {
+	t.Parallel()
+
+	externalModule := ast.MustParseModule(`package authz
+allow if input.user == "alice"
+deny if input.action == "delete"
+allowed if {
+	allow
+	not deny
+}`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	staticModule := ast.MustParseModule(`package main
+check if data.authz.allowed`)
+
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	input := ast.MustParseTerm(`{"user": "alice", "action": "read"}`)
+	qrs := runQuery(t, compiler, "data.main.check", input)
+
+	if len(qrs) != 1 {
+		t.Errorf("Expected 1 result, got %d", len(qrs))
+	}
+
+	if callCount := source.getCallCount(); callCount != 1 {
+		t.Errorf("Expected external source to be called once (cached for same ref and input), got %d calls", callCount)
+	}
+}
+
+type closableExternalSource struct {
+	refs       []ast.Ref
+	rules      []*ast.Rule
+	closeCalls int32
+}
+
+func (m *closableExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &closableExternalIndex{rules: m.rules, closeCalls: &m.closeCalls}, nil
+}
+
+func (m *closableExternalSource) Refs() []ast.Ref {
+	return m.refs
+}
+
+func (m *closableExternalSource) getCloseCalls() int {
+	return int(atomic.LoadInt32(&m.closeCalls))
+}
+
+type closableExternalIndex struct {
+	rules      []*ast.Rule
+	closeCalls *int32
+}
+
+func (*closableExternalIndex) Opts() *ast.ExternalSourceOptions {
+	return nil
+}
+
+func (m *closableExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	return m.rules, nil, nil
+}
+
+func (m *closableExternalIndex) Close() error {
+	atomic.AddInt32(m.closeCalls, 1)
+	return nil
+}
+
+func TestExternalSourceCloseCalled(t *testing.T) {
+	t.Parallel()
+
+	externalModule := ast.MustParseModule(`package authz
+allowed if input.user == "alice"`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &closableExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	staticModule := ast.MustParseModule(`package main
+check if data.authz.allowed`)
+
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	input := ast.MustParseTerm(`{"user": "alice"}`)
+	qrs := runQuery(t, compiler, "data.main.check", input)
+
+	if len(qrs) != 1 {
+		t.Errorf("Expected 1 result, got %d", len(qrs))
+	}
+
+	if closeCalls := source.getCloseCalls(); closeCalls != 1 {
+		t.Errorf("Expected Close() to be called once, got %d calls", closeCalls)
+	}
+}
+
+type preCompiledRulesSource struct {
+	refs          []ast.Ref
+	compiledRules []*ast.Rule
+}
+
+func (s *preCompiledRulesSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &preCompiledRulesIndex{compiledRules: s.compiledRules}, nil
+}
+
+func (s *preCompiledRulesSource) Refs() []ast.Ref {
+	return s.refs
+}
+
+type preCompiledRulesIndex struct {
+	compiledRules []*ast.Rule
+}
+
+func (*preCompiledRulesIndex) Opts() *ast.ExternalSourceOptions {
+	// For pre-compiled rules, skip all stages except those essential for
+	// integrating the rules into the compiler
+	var skippedStages []ast.StageID
+	essentialStages := []ast.StageID{
+		ast.StageSetModuleTree,
+		ast.StageSetRuleTree, ast.StageBuildRuleIndices,
+	}
+
+	for _, stage := range ast.AllStages() {
+		if !slices.Contains(essentialStages, stage) {
+			skippedStages = append(skippedStages, stage)
+		}
+	}
+
+	return &ast.ExternalSourceOptions{
+		SkippedStages: skippedStages,
+	}
+}
+
+func (idx *preCompiledRulesIndex) Lookup(_ context.Context, _ ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	return idx.compiledRules, nil, nil
+}
+
+func TestExternalSourceWithPreCompiledRules(t *testing.T) {
+	t.Parallel()
+
+	// Create and pre-compile an external module
+	externalModule := ast.MustParseModule(`package authz
+
+allow if input.user == "admin"
+deny if input.action == "delete"
+
+permitted if {
+	allow
+	not deny
+}`)
+
+	// Pre-compile the rules using a separate compiler
+	preCompiler := ast.NewCompiler()
+	preCompiler.Compile(map[string]*ast.Module{"authz.rego": externalModule})
+
+	if preCompiler.Failed() {
+		t.Fatalf("Pre-compilation failed: %v", preCompiler.Errors)
+	}
+
+	// Extract the pre-compiled rules
+	compiledRules := make([]*ast.Rule, 0, len(preCompiler.Modules))
+	for _, mod := range preCompiler.Modules {
+		compiledRules = append(compiledRules, mod.Rules...)
+	}
+
+	if len(compiledRules) == 0 {
+		t.Fatal("No compiled rules found")
+	}
+
+	// Create an external source that returns pre-compiled rules
+	packageRef := ast.MustParseRef("data.authz")
+	source := &preCompiledRulesSource{
+		refs:          []ast.Ref{packageRef},
+		compiledRules: compiledRules,
+	}
+
+	// Create a static module that uses the externally-provided rules
+	staticModule := ast.MustParseModule(`package main
+
+check if data.authz.permitted`)
+
+	// Set up compiler with the external source
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	t.Run("admin with read action should be allowed", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "admin", "action": "read"}`)
+		qrs := runQuery(t, compiler, "data.main.check", input)
+
+		if len(qrs) != 1 {
+			t.Errorf("Expected 1 result (allowed), got %d", len(qrs))
+		}
+	})
+
+	t.Run("admin with delete action should be denied", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "admin", "action": "delete"}`)
+		qrs := runQuery(t, compiler, "data.main.check", input)
+
+		if len(qrs) != 0 {
+			t.Errorf("Expected 0 results (denied), got %d", len(qrs))
+		}
+	})
+
+	t.Run("non-admin user should be denied", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "bob", "action": "read"}`)
+		qrs := runQuery(t, compiler, "data.main.check", input)
+
+		if len(qrs) != 0 {
+			t.Errorf("Expected 0 results (denied), got %d", len(qrs))
+		}
+	})
+}
+
+func TestExternalSourceDirectQuery(t *testing.T) {
+	t.Parallel()
+
+	externalModule := ast.MustParseModule(`package authz
+allow if input.user == "alice"
+deny if input.user == "bob"`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	compiler := setupCompiler(t, packageRef, source, nil)
+
+	t.Run("direct query allowed", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "alice"}`)
+		qrs := runQuery(t, compiler, "data.authz.allow", input)
+		if len(qrs) != 1 {
+			t.Errorf("Expected 1 result for direct external query, got %d", len(qrs))
+		}
+	})
+
+	t.Run("direct query denied", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "bob"}`)
+		qrs := runQuery(t, compiler, "data.authz.deny", input)
+		if len(qrs) != 1 {
+			t.Errorf("Expected 1 result for direct external query, got %d", len(qrs))
+		}
+	})
+
+	t.Run("direct query no match", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"user": "charlie"}`)
+		qrs := runQuery(t, compiler, "data.authz.allow", input)
+		if len(qrs) != 0 {
+			t.Errorf("Expected 0 results, got %d", len(qrs))
+		}
+	})
+}
+
+func TestExternalSourceCompilationFailure(t *testing.T) {
+	t.Parallel()
+
+	// Rules with an unsafe variable should fail compilation
+	externalModule := ast.MustParseModule(`package authz
+allow if { x }`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	staticModule := ast.MustParseModule(`package main
+check if data.authz.allow`)
+
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	store := inmem.New()
+	ctx := t.Context()
+	txn, err := store.NewTransaction(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Abort(ctx, txn)
+
+	query := ast.MustParseBody("data.main.check")
+	q := NewQuery(query).
+		WithCompiler(compiler).
+		WithStore(store).
+		WithTransaction(txn)
+
+	_, err = q.Run(ctx)
+	if err == nil {
+		t.Fatal("Expected compilation error from external source, got nil")
+	}
+	var errs ast.Errors
+	if !errors.As(err, &errs) {
+		t.Fatalf("Expected ast.Errors, got: %T: %v", err, err)
+	}
+	if !slices.ContainsFunc(errs, func(e *ast.Error) bool { return e.Code == ast.UnsafeVarErr }) {
+		t.Errorf("Expected unsafe var error, got: %v", err)
+	}
+}
+
+func TestExternalSourceE2EWithInputOverrideNilInput(t *testing.T) {
+	t.Parallel()
+
+	externalModule := ast.MustParseModule(`package authz
+allowed if input.user == "alice"`)
+
+	packageRef := ast.MustParseRef("data.authz")
+	source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules}
+
+	staticModule := ast.MustParseModule(`package main
+check if data.authz.allowed with input as {"user": "alice"}`)
+
+	compiler := setupCompiler(t, packageRef, source, staticModule)
+
+	// Query with nil input — the with clause provides it.
+	// This previously panicked because evalWithPop skipped PopFrame
+	// when oldInput was nil, leaking a frame on the externalTreeStack.
+	qrs := runQuery(t, compiler, "data.main.check", nil)
+
+	if len(qrs) != 1 {
+		t.Errorf("Expected 1 result, got %d", len(qrs))
+	}
+}

v1/topdown/instrumentation.go

@@ -19,6 +19,7 @@ const (
 	evalOpComprehensionCacheBuild = "eval_op_comprehension_cache_build"
 	evalOpComprehensionCacheHit   = "eval_op_comprehension_cache_hit"
 	evalOpComprehensionCacheMiss  = "eval_op_comprehension_cache_miss"
+	evalOpExternalRuleSource      = "eval_op_external_rule_source"
 	partialOpSaveUnify            = "partial_op_save_unify"
 	partialOpSaveSetContains      = "partial_op_save_set_contains"
 	partialOpSaveSetContainsRec   = "partial_op_save_set_contains_rec"

v1/topdown/save.go

@@ -357,7 +357,7 @@ func splitPackageAndRule(path ast.Ref) (ast.Ref, ast.Ref) {
 // being saved. This check allows the evaluator to evaluate statements
 // completely during partial evaluation as long as they do not depend on any
 // kind of unknown value or statements that would generate saves.
-func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, ss *saveSet, b *bindings, x any, rec bool) bool {
+func saveRequired(compilerTree *ast.TreeNode, extStack *externalTreeStack, ic *inliningControl, icIgnoreInternal bool, ss *saveSet, b *bindings, x any, rec bool) bool {
 
 	var found bool
 
@@ -389,8 +389,9 @@ func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, s
 				} else if ic.Disabled(v.ConstantPrefix(), icIgnoreInternal) {
 					found = true
 				} else {
-					for _, rule := range c.GetRulesDynamicWithOpts(v, ast.RulesOptions{IncludeHiddenModules: false}) {
-						if saveRequired(c, ic, icIgnoreInternal, ss, b, rule, true) {
+					rules := getRulesDynamic(compilerTree, extStack, v, ast.RulesOptions{IncludeHiddenModules: false})
+					for _, rule := range rules {
+						if saveRequired(compilerTree, extStack, ic, icIgnoreInternal, ss, b, rule, true) {
 							found = true
 							break
 						}
@@ -406,6 +407,76 @@ func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, s
 	return found
 }
 
+// getRulesDynamic looks up rules in both the compiler tree and external sources.
+func getRulesDynamic(compilerTree *ast.TreeNode, extStack *externalTreeStack, ref ast.Ref, opts ast.RulesOptions) []*ast.Rule {
+	var rules []*ast.Rule
+
+	// Check external trees
+	if extStack != nil {
+		for i := range extStack.entries {
+			entry := &extStack.entries[i]
+			if entry.tree != nil && ref.HasPrefix(entry.ref) {
+				// Navigate into the external tree using the remaining path
+				remaining := ref[len(entry.ref):]
+				rules = append(rules, getRulesFromTree(entry.tree, remaining, opts)...)
+			}
+		}
+	}
+
+	// Then check compiler tree
+	rules = append(rules, getRulesFromTree(compilerTree, ref, opts)...)
+
+	return rules
+}
+
+// getRulesFromTree walks a tree to find all rules matching the given ref.
+func getRulesFromTree(node *ast.TreeNode, ref ast.Ref, opts ast.RulesOptions) []*ast.Rule {
+	set := map[*ast.Rule]struct{}{}
+	var walk func(*ast.TreeNode, int)
+	walk = func(nav *ast.TreeNode, i int) {
+		switch {
+		case i >= len(ref):
+			nav.DepthFirst(func(descendant *ast.TreeNode) bool {
+				for _, rule := range descendant.Values {
+					set[rule] = struct{}{}
+				}
+				if opts.IncludeHiddenModules {
+					return false
+				}
+				return descendant.Hide
+			})
+
+		case i == 0 || ast.IsConstant(ref[i].Value):
+			if child := nav.Child(ref[i].Value); child != nil {
+				for _, rule := range child.Values {
+					set[rule] = struct{}{}
+				}
+				walk(child, i+1)
+			} else {
+				return
+			}
+
+		default:
+			for _, child := range nav.Children {
+				if child.Hide && !opts.IncludeHiddenModules {
+					continue
+				}
+				for _, rule := range child.Values {
+					set[rule] = struct{}{}
+				}
+				walk(child, i+1)
+			}
+		}
+	}
+
+	walk(node, 0)
+	rules := make([]*ast.Rule, 0, len(set))
+	for rule := range set {
+		rules = append(rules, rule)
+	}
+	return rules
+}
+
 func ignoreExprDuringPartial(expr *ast.Expr) bool {
 	if !expr.IsCall() {
 		return false