feat: composable EnableCondition with bitmask optimization

Add a composable EnableCondition system for tool filtering that:

1. **User-facing API** (conditions.go):
   - EnableCondition interface with Evaluate(ctx) method
   - Primitives: FeatureFlag(), ContextBool(), Static(), Always(), Never()
   - Combinators: And(), Or(), Not() with short-circuit evaluation
   - All bitmask complexity hidden from users

2. **Bitmask compiler** (condition_compiler.go):
   - Compiles conditions to O(1) bitmask evaluators at build time
   - RequestMask holds pre-computed uint64 bitmask per request
   - AND/OR of flags compile to single bitmask operations
   - Falls back gracefully for custom ConditionFunc

3. **Pre-sorting optimization** (builder.go):
   - Tools, resources, prompts sorted once at build time
   - Filtering preserves order, eliminating per-request sorting
   - ~45% faster request handling in benchmarks

4. **Integration** (filters.go, registry.go):
   - Builder.Build() compiles all EnableConditions
   - AvailableTools() builds RequestMask once, evaluates via bitmask
   - Backward compatible with legacy Enabled func and feature flags

Usage example:
  tool.EnableCondition = Or(
    ContextBool("is_cca"),           // CCA users bypass flag
    FeatureFlag("my_feature"),       // Others need flag enabled
  )

Benchmarks (1000 requests × 50 tools):
- Before: 23.7ms (with per-request sorting)
- After:  12.9ms (pre-sorted + bitmask)
- Improvement: 46% faster

This makes it easy for remote server to adopt - just set EnableCondition
on tools and the optimization is automatic.
This commit is contained in:
Sam Morrow
2025-12-18 14:20:35 +01:00
parent 3c453dd9dd
commit 3a52f7bd4b
10 changed files with 3429 additions and 51 deletions
+83 -4
View File
@@ -128,14 +128,22 @@ func (b *Builder) WithFilter(filter ToolFilter) *Builder {
}
// Build creates the final Inventory with all configuration applied.
// This processes toolset filtering, tool name resolution, and sets up
// This processes toolset filtering, tool name resolution, compiles EnableConditions
// for O(1) evaluation, pre-sorts all items for deterministic output, and sets up
// the inventory for use. The returned Inventory is ready for use with
// AvailableTools(), RegisterAll(), etc.
func (b *Builder) Build() *Inventory {
// Pre-sort tools, resources, and prompts at build time.
// This eliminates sorting overhead on every Available*() call.
// Filtering preserves order, so if input is sorted, output is sorted.
sortedTools := b.preSortTools()
sortedResources := b.preSortResources()
sortedPrompts := b.preSortPrompts()
r := &Inventory{
tools: b.tools,
resourceTemplates: b.resourceTemplates,
prompts: b.prompts,
tools: sortedTools,
resourceTemplates: sortedResources,
prompts: sortedPrompts,
deprecatedAliases: b.deprecatedAliases,
readOnly: b.readOnly,
featureChecker: b.featureChecker,
@@ -158,9 +166,80 @@ func (b *Builder) Build() *Inventory {
}
}
// Compile EnableConditions for O(1) bitmask evaluation
// Note: compileConditions uses r.tools which is now sortedTools
r.conditionCompiler, r.compiledConditions = b.compileConditions(sortedTools)
return r
}
// preSortTools returns a copy of tools sorted by toolset ID, then tool name.
// This allows filtering to preserve order without re-sorting.
func (b *Builder) preSortTools() []ServerTool {
if len(b.tools) == 0 {
return b.tools
}
sorted := make([]ServerTool, len(b.tools))
copy(sorted, b.tools)
sort.Slice(sorted, func(i, j int) bool {
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
}
return sorted[i].Tool.Name < sorted[j].Tool.Name
})
return sorted
}
// preSortResources returns a copy of resources sorted by toolset ID, then template name.
func (b *Builder) preSortResources() []ServerResourceTemplate {
if len(b.resourceTemplates) == 0 {
return b.resourceTemplates
}
sorted := make([]ServerResourceTemplate, len(b.resourceTemplates))
copy(sorted, b.resourceTemplates)
sort.Slice(sorted, func(i, j int) bool {
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
}
return sorted[i].Template.Name < sorted[j].Template.Name
})
return sorted
}
// preSortPrompts returns a copy of prompts sorted by toolset ID, then prompt name.
func (b *Builder) preSortPrompts() []ServerPrompt {
if len(b.prompts) == 0 {
return b.prompts
}
sorted := make([]ServerPrompt, len(b.prompts))
copy(sorted, b.prompts)
sort.Slice(sorted, func(i, j int) bool {
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
}
return sorted[i].Prompt.Name < sorted[j].Prompt.Name
})
return sorted
}
// compileConditions compiles all EnableConditions into bitmask-based evaluators.
// Returns the compiler (for building request masks) and compiled conditions slice.
// Takes the sorted tools slice to ensure compiled conditions align with sorted order.
func (b *Builder) compileConditions(sortedTools []ServerTool) (*ConditionCompiler, []*CompiledCondition) {
compiler := NewConditionCompiler()
compiled := make([]*CompiledCondition, len(sortedTools))
for i := range sortedTools {
if sortedTools[i].EnableCondition != nil {
compiled[i] = compiler.Compile(sortedTools[i].EnableCondition)
}
// nil means no condition (always enabled from condition perspective)
}
compiler.Freeze()
return compiler, compiled
}
// processToolsets processes the toolsetIDs configuration and returns:
// - enabledToolsets map (nil means all enabled)
// - unrecognizedToolsets list for warnings
+514
View File
@@ -0,0 +1,514 @@
package inventory
import (
"context"
"sync"
)
// ConditionCompiler compiles EnableConditions into optimized bitmask-based evaluators.
// This allows O(1) condition evaluation after an initial O(n) compilation phase.
//
// Design:
// 1. At build time, all tools register their EnableConditions with the compiler
// 2. The compiler analyzes conditions and assigns bit positions to each unique key
// 3. Each condition is compiled to a CompiledCondition with bitmask logic
// 4. At request time, all context bools are computed once into a RequestMask
// 5. Each tool's condition is evaluated via fast bitmask operations
//
// This trades memory (storing bit assignments) for speed (O(1) evaluation).
// For 50 tools with 10 unique condition keys, this saves ~40% evaluation time.
type ConditionCompiler struct {
mu sync.RWMutex
// keyToBit maps condition keys to bit positions (0-63)
// Keys are: "ctx:key_name" for ContextBool, "ff:flag_name" for FeatureFlag
keyToBit map[string]uint8
// nextBit is the next available bit position
nextBit uint8
// frozen prevents new bit assignments after compilation is complete
frozen bool
}
// NewConditionCompiler creates a new compiler for optimizing conditions.
func NewConditionCompiler() *ConditionCompiler {
return &ConditionCompiler{
keyToBit: make(map[string]uint8),
}
}
// assignBit returns the bit position for a key, assigning a new one if needed.
// Thread-safe. Panics if called after Freeze() and key doesn't exist.
func (cc *ConditionCompiler) assignBit(key string) uint8 {
cc.mu.RLock()
if bit, ok := cc.keyToBit[key]; ok {
cc.mu.RUnlock()
return bit
}
cc.mu.RUnlock()
cc.mu.Lock()
defer cc.mu.Unlock()
// Double-check after acquiring write lock
if bit, ok := cc.keyToBit[key]; ok {
return bit
}
if cc.frozen {
// After freezing, unknown keys get bit 63 (always false)
return 63
}
if cc.nextBit >= 63 {
// We've run out of bits - use bit 63 as overflow (always false)
return 63
}
bit := cc.nextBit
cc.keyToBit[key] = bit
cc.nextBit++
return bit
}
// Freeze prevents new bit assignments. Call after all conditions are compiled.
func (cc *ConditionCompiler) Freeze() {
cc.mu.Lock()
defer cc.mu.Unlock()
cc.frozen = true
}
// NumBits returns the number of bits assigned.
func (cc *ConditionCompiler) NumBits() int {
cc.mu.RLock()
defer cc.mu.RUnlock()
return int(cc.nextBit)
}
// Keys returns all registered keys (for debugging/introspection).
func (cc *ConditionCompiler) Keys() []string {
cc.mu.RLock()
defer cc.mu.RUnlock()
keys := make([]string, 0, len(cc.keyToBit))
for k := range cc.keyToBit {
keys = append(keys, k)
}
return keys
}
// Compile analyzes an EnableCondition and returns a CompiledCondition.
// The compiled condition uses bitmask operations for fast evaluation.
// Returns nil if the condition is nil (meaning always enabled).
func (cc *ConditionCompiler) Compile(cond EnableCondition) *CompiledCondition {
if cond == nil {
return nil // nil means always enabled
}
return cc.compile(cond)
}
func (cc *ConditionCompiler) compile(cond EnableCondition) *CompiledCondition {
switch c := cond.(type) {
case *staticCondition:
return &CompiledCondition{
evalType: evalStatic,
static: c.value,
}
case *contextBoolCondition:
bit := cc.assignBit("ctx:" + c.key)
return &CompiledCondition{
evalType: evalBitCheck,
requiredBit: bit,
requireTrue: true,
}
case *featureFlagCondition:
bit := cc.assignBit("ff:" + c.flagName)
return &CompiledCondition{
evalType: evalBitCheck,
requiredBit: bit,
requireTrue: true,
}
case *notCondition:
inner := cc.compile(c.condition)
if inner.evalType == evalStatic {
return &CompiledCondition{
evalType: evalStatic,
static: !inner.static,
}
}
if inner.evalType == evalBitCheck {
return &CompiledCondition{
evalType: evalBitCheck,
requiredBit: inner.requiredBit,
requireTrue: !inner.requireTrue,
}
}
return &CompiledCondition{
evalType: evalNot,
children: []*CompiledCondition{inner},
}
case *andCondition:
children := make([]*CompiledCondition, 0, len(c.conditions))
for _, child := range c.conditions {
compiled := cc.compile(child)
// Optimize: static false short-circuits entire AND
if compiled.evalType == evalStatic && !compiled.static {
return &CompiledCondition{evalType: evalStatic, static: false}
}
// Optimize: skip static true (no-op in AND)
if compiled.evalType == evalStatic && compiled.static {
continue
}
children = append(children, compiled)
}
if len(children) == 0 {
return &CompiledCondition{evalType: evalStatic, static: true}
}
if len(children) == 1 {
return children[0]
}
// Check if we can use bitmask AND (all children are simple bit checks with requireTrue)
if canUseBitmaskAnd(children) {
var mask uint64
for _, child := range children {
mask |= 1 << child.requiredBit
}
return &CompiledCondition{
evalType: evalBitmaskAnd,
bitmask: mask,
requireTrue: true,
}
}
return &CompiledCondition{
evalType: evalAnd,
children: children,
}
case *orCondition:
children := make([]*CompiledCondition, 0, len(c.conditions))
for _, child := range c.conditions {
compiled := cc.compile(child)
// Optimize: static true short-circuits entire OR
if compiled.evalType == evalStatic && compiled.static {
return &CompiledCondition{evalType: evalStatic, static: true}
}
// Optimize: skip static false (no-op in OR)
if compiled.evalType == evalStatic && !compiled.static {
continue
}
children = append(children, compiled)
}
if len(children) == 0 {
return &CompiledCondition{evalType: evalStatic, static: false}
}
if len(children) == 1 {
return children[0]
}
// Check if we can use bitmask OR (all children are simple bit checks with requireTrue)
if canUseBitmaskOr(children) {
var mask uint64
for _, child := range children {
mask |= 1 << child.requiredBit
}
return &CompiledCondition{
evalType: evalBitmaskOr,
bitmask: mask,
requireTrue: true,
}
}
return &CompiledCondition{
evalType: evalOr,
children: children,
}
case ConditionFunc:
// Can't optimize arbitrary functions - fall back to direct evaluation
return &CompiledCondition{
evalType: evalFallback,
fallback: c,
}
default:
// Unknown condition type - fall back to direct evaluation
return &CompiledCondition{
evalType: evalFallback,
fallback: cond,
}
}
}
// canUseBitmaskAnd checks if all children are simple positive bit checks
func canUseBitmaskAnd(children []*CompiledCondition) bool {
for _, c := range children {
if c.evalType != evalBitCheck || !c.requireTrue {
return false
}
}
return true
}
// canUseBitmaskOr checks if all children are simple positive bit checks
func canUseBitmaskOr(children []*CompiledCondition) bool {
for _, c := range children {
if c.evalType != evalBitCheck || !c.requireTrue {
return false
}
}
return true
}
// evalType describes how a CompiledCondition should be evaluated
type evalType uint8
const (
evalStatic evalType = iota // Return static value
evalBitCheck // Check single bit
evalBitmaskAnd // AND: (mask & bits) == mask
evalBitmaskOr // OR: (mask & bits) != 0
evalAnd // Tree-based AND
evalOr // Tree-based OR
evalNot // Negate child
evalFallback // Call original condition
)
// CompiledCondition is an optimized representation of an EnableCondition.
// It uses bitmask operations where possible for O(1) evaluation.
type CompiledCondition struct {
evalType evalType
// For evalStatic
static bool
// For evalBitCheck
requiredBit uint8
requireTrue bool // true = bit must be set, false = bit must be unset
// For evalBitmaskAnd/evalBitmaskOr
bitmask uint64
// For evalAnd/evalOr/evalNot
children []*CompiledCondition
// For evalFallback
fallback EnableCondition
}
// Evaluate checks the compiled condition against the given request mask.
// For most conditions this is O(1) - just bitmask operations.
func (cc *CompiledCondition) Evaluate(rm *RequestMask) (bool, error) {
switch cc.evalType {
case evalStatic:
return cc.static, nil
case evalBitCheck:
bitSet := (rm.bits & (1 << cc.requiredBit)) != 0
if cc.requireTrue {
return bitSet, nil
}
return !bitSet, nil
case evalBitmaskAnd:
// All required bits must be set
return (rm.bits & cc.bitmask) == cc.bitmask, nil
case evalBitmaskOr:
// Any required bit must be set
return (rm.bits & cc.bitmask) != 0, nil
case evalAnd:
for _, child := range cc.children {
result, err := child.Evaluate(rm)
if err != nil {
return false, err
}
if !result {
return false, nil
}
}
return true, nil
case evalOr:
for _, child := range cc.children {
result, err := child.Evaluate(rm)
if err != nil {
continue // OR continues on error
}
if result {
return true, nil
}
}
return false, nil
case evalNot:
result, err := cc.children[0].Evaluate(rm)
if err != nil {
return false, err
}
return !result, nil
case evalFallback:
return cc.fallback.Evaluate(rm.ctx)
default:
return false, nil
}
}
// RequestMask holds pre-computed condition values as a bitmask.
// Created once per request, then used to evaluate all tool conditions.
type RequestMask struct {
bits uint64
ctx context.Context // For fallback evaluation
}
// RequestMaskBuilder builds a RequestMask from context bools and feature flags.
type RequestMaskBuilder struct {
compiler *ConditionCompiler
}
// NewRequestMaskBuilder creates a builder for the given compiler.
func NewRequestMaskBuilder(compiler *ConditionCompiler) *RequestMaskBuilder {
return &RequestMaskBuilder{compiler: compiler}
}
// Build creates a RequestMask from context bools and feature flag results.
// This should be called once per request with all relevant bools pre-computed.
func (b *RequestMaskBuilder) Build(ctx context.Context, bools ContextBools, flags map[string]bool) *RequestMask {
var bits uint64
b.compiler.mu.RLock()
defer b.compiler.mu.RUnlock()
// Set bits for context bools
for key, value := range bools {
if bit, ok := b.compiler.keyToBit["ctx:"+key]; ok && value {
bits |= 1 << bit
}
}
// Set bits for feature flags
for flag, enabled := range flags {
if bit, ok := b.compiler.keyToBit["ff:"+flag]; ok && enabled {
bits |= 1 << bit
}
}
return &RequestMask{
bits: bits,
ctx: ctx,
}
}
// BuildFromContext creates a RequestMask using ContextBools from context
// and evaluating feature flags via the FeatureFlagChecker in context.
// This is a convenience method that computes everything from context.
func (b *RequestMaskBuilder) BuildFromContext(ctx context.Context) *RequestMask {
var bits uint64
bools := contextBoolsFromContext(ctx)
checker := FeatureCheckerFromContext(ctx)
b.compiler.mu.RLock()
defer b.compiler.mu.RUnlock()
for key, bit := range b.compiler.keyToBit {
if len(key) < 4 {
continue
}
prefix := key[:3]
name := key[3:]
switch prefix {
case "ctx":
if bools != nil && bools[name] {
bits |= 1 << bit
}
case "ff:":
name = key[3:] // "ff:" is 3 chars
if checker != nil {
enabled, err := checker(ctx, name)
if err == nil && enabled {
bits |= 1 << bit
}
}
}
}
return &RequestMask{
bits: bits,
ctx: ctx,
}
}
// --- Integration with Inventory ---
// CompiledToolCondition pairs a tool with its compiled condition.
type CompiledToolCondition struct {
Tool *ServerTool
Condition *CompiledCondition // nil means always enabled
}
// ToolConditionSet holds all compiled tool conditions for fast filtering.
type ToolConditionSet struct {
compiler *ConditionCompiler
builder *RequestMaskBuilder
tools []CompiledToolCondition
}
// NewToolConditionSet creates a new set from the given tools.
// This compiles all conditions and freezes the compiler.
func NewToolConditionSet(tools []*ServerTool) *ToolConditionSet {
compiler := NewConditionCompiler()
compiled := make([]CompiledToolCondition, len(tools))
for i, tool := range tools {
compiled[i] = CompiledToolCondition{
Tool: tool,
Condition: compiler.Compile(tool.EnableCondition),
}
}
compiler.Freeze()
return &ToolConditionSet{
compiler: compiler,
builder: NewRequestMaskBuilder(compiler),
tools: compiled,
}
}
// FilterEnabled returns tools that are enabled for the given request mask.
func (tcs *ToolConditionSet) FilterEnabled(rm *RequestMask) []*ServerTool {
result := make([]*ServerTool, 0, len(tcs.tools))
for _, tc := range tcs.tools {
if tc.Condition == nil {
// No condition = always enabled
result = append(result, tc.Tool)
continue
}
enabled, _ := tc.Condition.Evaluate(rm)
if enabled {
result = append(result, tc.Tool)
}
}
return result
}
// BuildMask creates a RequestMask for filtering.
func (tcs *ToolConditionSet) BuildMask(ctx context.Context, bools ContextBools, flags map[string]bool) *RequestMask {
return tcs.builder.Build(ctx, bools, flags)
}
// BuildMaskFromContext creates a RequestMask from context.
func (tcs *ToolConditionSet) BuildMaskFromContext(ctx context.Context) *RequestMask {
return tcs.builder.BuildFromContext(ctx)
}
// Compiler returns the condition compiler (for introspection/debugging).
func (tcs *ToolConditionSet) Compiler() *ConditionCompiler {
return tcs.compiler
}
+442
View File
@@ -0,0 +1,442 @@
package inventory
import (
"context"
"testing"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestConditionCompiler_AssignBit(t *testing.T) {
cc := NewConditionCompiler()
// First assignment
bit1 := cc.assignBit("ctx:is_cca")
assert.Equal(t, uint8(0), bit1)
// Same key returns same bit
bit2 := cc.assignBit("ctx:is_cca")
assert.Equal(t, bit1, bit2)
// Different key gets new bit
bit3 := cc.assignBit("ff:web_search")
assert.Equal(t, uint8(1), bit3)
// After freeze, unknown keys get bit 63
cc.Freeze()
bit4 := cc.assignBit("ctx:unknown")
assert.Equal(t, uint8(63), bit4)
// Known keys still work after freeze
bit5 := cc.assignBit("ctx:is_cca")
assert.Equal(t, bit1, bit5)
}
func TestConditionCompiler_CompileStatic(t *testing.T) {
cc := NewConditionCompiler()
// Static true
cond := cc.Compile(Always())
require.NotNil(t, cond)
assert.Equal(t, evalStatic, cond.evalType)
assert.True(t, cond.static)
// Static false
cond = cc.Compile(Never())
require.NotNil(t, cond)
assert.Equal(t, evalStatic, cond.evalType)
assert.False(t, cond.static)
// Nil returns nil
assert.Nil(t, cc.Compile(nil))
}
func TestConditionCompiler_CompileContextBool(t *testing.T) {
cc := NewConditionCompiler()
cond := cc.Compile(ContextBool("is_cca"))
require.NotNil(t, cond)
assert.Equal(t, evalBitCheck, cond.evalType)
assert.Equal(t, uint8(0), cond.requiredBit)
assert.True(t, cond.requireTrue)
}
func TestConditionCompiler_CompileFeatureFlag(t *testing.T) {
cc := NewConditionCompiler()
cond := cc.Compile(FeatureFlag("web_search"))
require.NotNil(t, cond)
assert.Equal(t, evalBitCheck, cond.evalType)
assert.Equal(t, uint8(0), cond.requiredBit)
assert.True(t, cond.requireTrue)
}
func TestConditionCompiler_CompileNot(t *testing.T) {
cc := NewConditionCompiler()
// Not(static) -> static
cond := cc.Compile(Not(Always()))
require.NotNil(t, cond)
assert.Equal(t, evalStatic, cond.evalType)
assert.False(t, cond.static)
// Not(contextBool) -> bitCheck with requireTrue=false
cond = cc.Compile(Not(ContextBool("is_cca")))
require.NotNil(t, cond)
assert.Equal(t, evalBitCheck, cond.evalType)
assert.False(t, cond.requireTrue)
}
func TestConditionCompiler_CompileAnd(t *testing.T) {
cc := NewConditionCompiler()
// And of two context bools -> bitmaskAnd
cond := cc.Compile(And(
ContextBool("is_cca"),
ContextBool("has_access"),
))
require.NotNil(t, cond)
assert.Equal(t, evalBitmaskAnd, cond.evalType)
assert.Equal(t, uint64(0b11), cond.bitmask) // bits 0 and 1
// And with static false -> static false
cond = cc.Compile(And(
ContextBool("is_cca"),
Never(),
))
require.NotNil(t, cond)
assert.Equal(t, evalStatic, cond.evalType)
assert.False(t, cond.static)
// And with static true filtered out -> single condition
cc2 := NewConditionCompiler()
cond = cc2.Compile(And(
Always(),
ContextBool("is_cca"),
))
require.NotNil(t, cond)
assert.Equal(t, evalBitCheck, cond.evalType)
}
func TestConditionCompiler_CompileOr(t *testing.T) {
cc := NewConditionCompiler()
// Or of two context bools -> bitmaskOr
cond := cc.Compile(Or(
ContextBool("is_cca"),
ContextBool("is_bypass"),
))
require.NotNil(t, cond)
assert.Equal(t, evalBitmaskOr, cond.evalType)
assert.Equal(t, uint64(0b11), cond.bitmask)
// Or with static true -> static true
cond = cc.Compile(Or(
ContextBool("is_cca"),
Always(),
))
require.NotNil(t, cond)
assert.Equal(t, evalStatic, cond.evalType)
assert.True(t, cond.static)
}
func TestCompiledCondition_Evaluate(t *testing.T) {
cc := NewConditionCompiler()
// Compile conditions
ccaCheck := cc.Compile(ContextBool("is_cca"))
ffCheck := cc.Compile(FeatureFlag("web_search"))
andCond := cc.Compile(And(ContextBool("is_cca"), FeatureFlag("web_search")))
orCond := cc.Compile(Or(ContextBool("is_cca"), FeatureFlag("web_search")))
notCond := cc.Compile(Not(ContextBool("is_cca")))
cc.Freeze()
builder := NewRequestMaskBuilder(cc)
tests := []struct {
name string
condition *CompiledCondition
bools ContextBools
flags map[string]bool
want bool
}{
{
name: "context bool true",
condition: ccaCheck,
bools: ContextBools{"is_cca": true},
want: true,
},
{
name: "context bool false",
condition: ccaCheck,
bools: ContextBools{"is_cca": false},
want: false,
},
{
name: "feature flag true",
condition: ffCheck,
flags: map[string]bool{"web_search": true},
want: true,
},
{
name: "feature flag false",
condition: ffCheck,
flags: map[string]bool{"web_search": false},
want: false,
},
{
name: "and both true",
condition: andCond,
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"web_search": true},
want: true,
},
{
name: "and one false",
condition: andCond,
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"web_search": false},
want: false,
},
{
name: "or one true",
condition: orCond,
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"web_search": false},
want: true,
},
{
name: "or both false",
condition: orCond,
bools: ContextBools{"is_cca": false},
flags: map[string]bool{"web_search": false},
want: false,
},
{
name: "not true -> false",
condition: notCond,
bools: ContextBools{"is_cca": true},
want: false,
},
{
name: "not false -> true",
condition: notCond,
bools: ContextBools{"is_cca": false},
want: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mask := builder.Build(context.Background(), tt.bools, tt.flags)
got, err := tt.condition.Evaluate(mask)
require.NoError(t, err)
assert.Equal(t, tt.want, got)
})
}
}
func TestToolConditionSet_FilterEnabled(t *testing.T) {
tools := []*ServerTool{
{
Tool: mcp.Tool{Name: "always_on"},
EnableCondition: nil, // nil means always enabled
},
{
Tool: mcp.Tool{Name: "cca_only"},
EnableCondition: ContextBool("is_cca"),
},
{
Tool: mcp.Tool{Name: "ff_required"},
EnableCondition: FeatureFlag("web_search"),
},
{
Tool: mcp.Tool{Name: "cca_and_ff"},
EnableCondition: And(
ContextBool("is_cca"),
FeatureFlag("code_search"),
),
},
{
Tool: mcp.Tool{Name: "cca_or_ff"},
EnableCondition: Or(
ContextBool("is_cca"),
FeatureFlag("bypass_flag"),
),
},
}
tcs := NewToolConditionSet(tools)
tests := []struct {
name string
bools ContextBools
flags map[string]bool
want []string
}{
{
name: "no bools or flags - only always_on",
bools: nil,
flags: nil,
want: []string{"always_on"},
},
{
name: "cca only",
bools: ContextBools{"is_cca": true},
flags: nil,
want: []string{"always_on", "cca_only", "cca_or_ff"},
},
{
name: "web_search flag only",
bools: nil,
flags: map[string]bool{"web_search": true},
want: []string{"always_on", "ff_required"},
},
{
name: "cca and code_search",
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"code_search": true},
want: []string{"always_on", "cca_only", "cca_and_ff", "cca_or_ff"},
},
{
name: "all enabled",
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"web_search": true, "code_search": true, "bypass_flag": true},
want: []string{"always_on", "cca_only", "ff_required", "cca_and_ff", "cca_or_ff"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mask := tcs.BuildMask(context.Background(), tt.bools, tt.flags)
enabled := tcs.FilterEnabled(mask)
names := make([]string, len(enabled))
for i, tool := range enabled {
names[i] = tool.Tool.Name
}
assert.Equal(t, tt.want, names)
})
}
}
func TestToolConditionSet_ComplexConditions(t *testing.T) {
// Test complex real-world patterns from remote server
tools := []*ServerTool{
{
// CCA bypass: CCA OR feature_flag
Tool: mcp.Tool{Name: "cca_bypass"},
EnableCondition: Or(
ContextBool("is_cca"),
FeatureFlag("agent_search"),
),
},
{
// Feature + policy: feature AND paid_access
Tool: mcp.Tool{Name: "paid_feature"},
EnableCondition: And(
FeatureFlag("premium_search"),
ContextBool("has_paid_access"),
),
},
{
// Complex: (CCA OR copilot_chat) AND feature AND NOT disabled
Tool: mcp.Tool{Name: "complex"},
EnableCondition: And(
Or(
ContextBool("is_cca"),
ContextBool("is_copilot_chat"),
),
FeatureFlag("advanced_feature"),
Not(FeatureFlag("kill_switch")),
),
},
}
tcs := NewToolConditionSet(tools)
tests := []struct {
name string
bools ContextBools
flags map[string]bool
want []string
}{
{
name: "cca enables cca_bypass",
bools: ContextBools{"is_cca": true},
flags: nil,
want: []string{"cca_bypass"},
},
{
name: "agent_search flag enables cca_bypass",
bools: nil,
flags: map[string]bool{"agent_search": true},
want: []string{"cca_bypass"},
},
{
name: "premium + paid enables paid_feature",
bools: ContextBools{"has_paid_access": true},
flags: map[string]bool{"premium_search": true},
want: []string{"paid_feature"},
},
{
name: "complex enabled with cca + feature",
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"advanced_feature": true},
want: []string{"cca_bypass", "complex"},
},
{
name: "complex disabled by kill_switch",
bools: ContextBools{"is_cca": true},
flags: map[string]bool{"advanced_feature": true, "kill_switch": true},
want: []string{"cca_bypass"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mask := tcs.BuildMask(context.Background(), tt.bools, tt.flags)
enabled := tcs.FilterEnabled(mask)
names := make([]string, len(enabled))
for i, tool := range enabled {
names[i] = tool.Tool.Name
}
assert.Equal(t, tt.want, names)
})
}
}
func TestConditionCompiler_NumBits(t *testing.T) {
cc := NewConditionCompiler()
assert.Equal(t, 0, cc.NumBits())
cc.assignBit("ctx:a")
assert.Equal(t, 1, cc.NumBits())
cc.assignBit("ctx:b")
cc.assignBit("ff:c")
assert.Equal(t, 3, cc.NumBits())
// Same key doesn't increase count
cc.assignBit("ctx:a")
assert.Equal(t, 3, cc.NumBits())
}
func TestConditionCompiler_Keys(t *testing.T) {
cc := NewConditionCompiler()
cc.assignBit("ctx:is_cca")
cc.assignBit("ff:web_search")
cc.assignBit("ctx:has_access")
keys := cc.Keys()
assert.Len(t, keys, 3)
assert.Contains(t, keys, "ctx:is_cca")
assert.Contains(t, keys, "ff:web_search")
assert.Contains(t, keys, "ctx:has_access")
}
+293
View File
@@ -0,0 +1,293 @@
package inventory
import (
"context"
)
// EnableCondition represents a composable condition for tool availability.
// Conditions can be combined using And/Or/Not combinators for complex logic.
//
// Design goals:
// - Declarative: users compose conditions without knowing implementation details
// - Composable: complex conditions built from simple primitives
// - Efficient: conditions are evaluated lazily, with results potentially cached
// - Decoupled: condition definitions don't depend on specific actor types
//
// Example usage:
//
// // Simple feature flag
// tool.EnableCondition = FeatureFlag("web_search")
//
// // Feature flag AND user policy
// tool.EnableCondition = And(
// FeatureFlag("web_search"),
// ContextBool("user_has_paid_bing_access"),
// )
//
// // CCA bypass (CCA requests OR feature flag for non-CCA)
// tool.EnableCondition = Or(
// ContextBool("is_cca"),
// FeatureFlag("agent_search"),
// )
type EnableCondition interface {
// Evaluate checks if the condition is met in the given context.
// Returns (enabled, error). On error, the condition should be treated as false.
Evaluate(ctx context.Context) (bool, error)
}
// ConditionFunc is an adapter that allows functions to be used as EnableConditions.
// This is useful for simple one-off conditions that don't need to be reusable.
type ConditionFunc func(ctx context.Context) (bool, error)
// Evaluate implements EnableCondition.
func (f ConditionFunc) Evaluate(ctx context.Context) (bool, error) {
return f(ctx)
}
// --- Primitive Conditions ---
// featureFlagCondition checks if a named feature flag is enabled.
// The actual flag checking is delegated to a FeatureFlagChecker in context.
type featureFlagCondition struct {
flagName string
}
// FeatureFlag creates a condition that checks if the named feature flag is enabled.
// The feature flag is evaluated using the FeatureFlagChecker stored in context.
// If no checker is available or if the flag check returns an error, the condition is false.
func FeatureFlag(flagName string) EnableCondition {
return &featureFlagCondition{flagName: flagName}
}
// Evaluate implements EnableCondition.
func (c *featureFlagCondition) Evaluate(ctx context.Context) (bool, error) {
checker := FeatureCheckerFromContext(ctx)
if checker == nil {
return false, nil
}
return checker(ctx, c.flagName)
}
// contextBoolCondition checks a named boolean value from context.
// This allows tools to depend on pre-computed boolean conditions without
// knowing how those conditions are computed.
type contextBoolCondition struct {
key string
}
// ContextBool creates a condition that checks a named boolean from context.
// The boolean is retrieved using ContextBoolFromContext(ctx, key).
// This decouples tool definitions from specific actor/user types.
//
// Common keys might include:
// - "is_cca" - whether this is a Copilot Coding Agent request
// - "user_has_paid_access" - whether user has paid Copilot access
// - "mcp_host_is_copilot_chat" - whether MCP host is copilot-chat
//
// Returns false if the key is not found in context.
func ContextBool(key string) EnableCondition {
return &contextBoolCondition{key: key}
}
// Evaluate implements EnableCondition.
func (c *contextBoolCondition) Evaluate(ctx context.Context) (bool, error) {
return ContextBoolFromContext(ctx, c.key), nil
}
// staticCondition always returns a fixed value.
type staticCondition struct {
value bool
}
// Static creates a condition that always returns the given value.
// Useful for testing or for conditions that are determined at build time.
func Static(value bool) EnableCondition {
return &staticCondition{value: value}
}
// Always returns a condition that is always true.
// Useful as a default or placeholder.
func Always() EnableCondition {
return Static(true)
}
// Never returns a condition that is always false.
// Useful for disabling tools unconditionally.
func Never() EnableCondition {
return Static(false)
}
// Evaluate implements EnableCondition.
func (c *staticCondition) Evaluate(_ context.Context) (bool, error) {
return c.value, nil
}
// --- Combinators ---
// andCondition requires all conditions to be true.
type andCondition struct {
conditions []EnableCondition
}
// And creates a condition that is true only if ALL of the given conditions are true.
// Short-circuits on the first false condition.
// Returns true if no conditions are provided.
func And(conditions ...EnableCondition) EnableCondition {
// Filter out nil conditions
filtered := make([]EnableCondition, 0, len(conditions))
for _, c := range conditions {
if c != nil {
filtered = append(filtered, c)
}
}
if len(filtered) == 0 {
return Always()
}
if len(filtered) == 1 {
return filtered[0]
}
return &andCondition{conditions: filtered}
}
// Evaluate implements EnableCondition.
func (c *andCondition) Evaluate(ctx context.Context) (bool, error) {
for _, cond := range c.conditions {
enabled, err := cond.Evaluate(ctx)
if err != nil {
return false, err
}
if !enabled {
return false, nil
}
}
return true, nil
}
// orCondition requires at least one condition to be true.
type orCondition struct {
conditions []EnableCondition
}
// Or creates a condition that is true if ANY of the given conditions is true.
// Short-circuits on the first true condition.
// Returns false if no conditions are provided.
func Or(conditions ...EnableCondition) EnableCondition {
// Filter out nil conditions
filtered := make([]EnableCondition, 0, len(conditions))
for _, c := range conditions {
if c != nil {
filtered = append(filtered, c)
}
}
if len(filtered) == 0 {
return Never()
}
if len(filtered) == 1 {
return filtered[0]
}
return &orCondition{conditions: filtered}
}
// Evaluate implements EnableCondition.
func (c *orCondition) Evaluate(ctx context.Context) (bool, error) {
for _, cond := range c.conditions {
enabled, err := cond.Evaluate(ctx)
if err != nil {
// For OR, we continue checking other conditions on error
continue
}
if enabled {
return true, nil
}
}
return false, nil
}
// notCondition negates a condition.
type notCondition struct {
condition EnableCondition
}
// Not creates a condition that is the logical negation of the given condition.
// Returns true if the inner condition returns false (and vice versa).
// Errors are propagated.
func Not(condition EnableCondition) EnableCondition {
if condition == nil {
return Never() // Not(nil) = Not(true) = false
}
return &notCondition{condition: condition}
}
// Evaluate implements EnableCondition.
func (c *notCondition) Evaluate(ctx context.Context) (bool, error) {
enabled, err := c.condition.Evaluate(ctx)
if err != nil {
return false, err
}
return !enabled, nil
}
// --- Context Keys for Conditions ---
// Context key types for storing condition-related data
type contextKey int
const (
featureCheckerKey contextKey = iota
contextBoolsKey
)
// ContextWithFeatureChecker returns a context with the given feature flag checker.
func ContextWithFeatureChecker(ctx context.Context, checker FeatureFlagChecker) context.Context {
return context.WithValue(ctx, featureCheckerKey, checker)
}
// FeatureCheckerFromContext retrieves the feature flag checker from context.
// Returns nil if no checker is set.
func FeatureCheckerFromContext(ctx context.Context) FeatureFlagChecker {
checker, _ := ctx.Value(featureCheckerKey).(FeatureFlagChecker)
return checker
}
// ContextBools is a map of named boolean values for use with ContextBool conditions.
// This allows callers to pre-compute common checks once per request and share them.
type ContextBools map[string]bool
// ContextWithBools returns a context with the given boolean values.
// These values can be retrieved using ContextBool conditions.
func ContextWithBools(ctx context.Context, bools ContextBools) context.Context {
// Merge with existing bools if any
existing := contextBoolsFromContext(ctx)
if existing != nil {
merged := make(ContextBools, len(existing)+len(bools))
for k, v := range existing {
merged[k] = v
}
for k, v := range bools {
merged[k] = v
}
return context.WithValue(ctx, contextBoolsKey, merged)
}
return context.WithValue(ctx, contextBoolsKey, bools)
}
// contextBoolsFromContext retrieves all context bools.
func contextBoolsFromContext(ctx context.Context) ContextBools {
bools, _ := ctx.Value(contextBoolsKey).(ContextBools)
return bools
}
// ContextBoolFromContext retrieves a named boolean from context.
// Returns false if the key is not found.
func ContextBoolFromContext(ctx context.Context, key string) bool {
bools := contextBoolsFromContext(ctx)
if bools == nil {
return false
}
return bools[key]
}
// SetContextBool is a convenience function that adds a single boolean to context.
func SetContextBool(ctx context.Context, key string, value bool) context.Context {
return ContextWithBools(ctx, ContextBools{key: value})
}
+552
View File
@@ -0,0 +1,552 @@
package inventory
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestFeatureFlagCondition(t *testing.T) {
tests := []struct {
name string
flagName string
checkerResult bool
checkerErr error
hasChecker bool
expectedResult bool
expectedErr bool
}{
{
name: "flag enabled",
flagName: "test_flag",
checkerResult: true,
hasChecker: true,
expectedResult: true,
},
{
name: "flag disabled",
flagName: "test_flag",
checkerResult: false,
hasChecker: true,
expectedResult: false,
},
{
name: "no checker in context",
flagName: "test_flag",
hasChecker: false,
expectedResult: false,
},
{
name: "checker returns error",
flagName: "test_flag",
checkerErr: errors.New("flag check failed"),
hasChecker: true,
expectedResult: false,
expectedErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
if tt.hasChecker {
checker := func(_ context.Context, flagName string) (bool, error) {
assert.Equal(t, tt.flagName, flagName)
return tt.checkerResult, tt.checkerErr
}
ctx = ContextWithFeatureChecker(ctx, checker)
}
cond := FeatureFlag(tt.flagName)
result, err := cond.Evaluate(ctx)
if tt.expectedErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
assert.Equal(t, tt.expectedResult, result)
})
}
}
func TestContextBoolCondition(t *testing.T) {
tests := []struct {
name string
key string
contextBools ContextBools
expectedResult bool
}{
{
name: "bool is true",
key: "is_cca",
contextBools: ContextBools{"is_cca": true},
expectedResult: true,
},
{
name: "bool is false",
key: "is_cca",
contextBools: ContextBools{"is_cca": false},
expectedResult: false,
},
{
name: "key not found",
key: "is_cca",
contextBools: ContextBools{"other_key": true},
expectedResult: false,
},
{
name: "no context bools",
key: "is_cca",
contextBools: nil,
expectedResult: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
if tt.contextBools != nil {
ctx = ContextWithBools(ctx, tt.contextBools)
}
cond := ContextBool(tt.key)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.Equal(t, tt.expectedResult, result)
})
}
}
func TestStaticConditions(t *testing.T) {
ctx := context.Background()
t.Run("Static(true)", func(t *testing.T) {
cond := Static(true)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
})
t.Run("Static(false)", func(t *testing.T) {
cond := Static(false)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.False(t, result)
})
t.Run("Always()", func(t *testing.T) {
cond := Always()
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
})
t.Run("Never()", func(t *testing.T) {
cond := Never()
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.False(t, result)
})
}
func TestAndCondition(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
conditions []EnableCondition
expectedResult bool
}{
{
name: "all true",
conditions: []EnableCondition{Always(), Always(), Always()},
expectedResult: true,
},
{
name: "one false",
conditions: []EnableCondition{Always(), Never(), Always()},
expectedResult: false,
},
{
name: "all false",
conditions: []EnableCondition{Never(), Never()},
expectedResult: false,
},
{
name: "empty conditions",
conditions: []EnableCondition{},
expectedResult: true,
},
{
name: "single true",
conditions: []EnableCondition{Always()},
expectedResult: true,
},
{
name: "single false",
conditions: []EnableCondition{Never()},
expectedResult: false,
},
{
name: "nil conditions filtered out",
conditions: []EnableCondition{Always(), nil, Always()},
expectedResult: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cond := And(tt.conditions...)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.Equal(t, tt.expectedResult, result)
})
}
}
func TestOrCondition(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
conditions []EnableCondition
expectedResult bool
}{
{
name: "all true",
conditions: []EnableCondition{Always(), Always(), Always()},
expectedResult: true,
},
{
name: "one true",
conditions: []EnableCondition{Never(), Always(), Never()},
expectedResult: true,
},
{
name: "all false",
conditions: []EnableCondition{Never(), Never()},
expectedResult: false,
},
{
name: "empty conditions",
conditions: []EnableCondition{},
expectedResult: false,
},
{
name: "single true",
conditions: []EnableCondition{Always()},
expectedResult: true,
},
{
name: "single false",
conditions: []EnableCondition{Never()},
expectedResult: false,
},
{
name: "nil conditions filtered out",
conditions: []EnableCondition{Never(), nil, Never()},
expectedResult: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cond := Or(tt.conditions...)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.Equal(t, tt.expectedResult, result)
})
}
}
func TestNotCondition(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
condition EnableCondition
expectedResult bool
}{
{
name: "not true",
condition: Always(),
expectedResult: false,
},
{
name: "not false",
condition: Never(),
expectedResult: true,
},
{
name: "not nil",
condition: nil,
expectedResult: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cond := Not(tt.condition)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.Equal(t, tt.expectedResult, result)
})
}
}
func TestConditionFunc(t *testing.T) {
ctx := context.Background()
t.Run("simple function", func(t *testing.T) {
cond := ConditionFunc(func(_ context.Context) (bool, error) {
return true, nil
})
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
})
t.Run("function with error", func(t *testing.T) {
expectedErr := errors.New("test error")
cond := ConditionFunc(func(_ context.Context) (bool, error) {
return false, expectedErr
})
result, err := cond.Evaluate(ctx)
assert.Equal(t, expectedErr, err)
assert.False(t, result)
})
}
func TestComplexConditionCombinations(t *testing.T) {
// These tests match the real-world scenarios from the remote server
t.Run("feature flag AND user policy (web search pattern)", func(t *testing.T) {
// Pattern: feature flag must be enabled AND user must have paid Bing access
cond := And(
FeatureFlag("web_search"),
ContextBool("user_has_paid_bing_access"),
)
// Test: both conditions true
ctx := context.Background()
ctx = ContextWithFeatureChecker(ctx, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx = ContextWithBools(ctx, ContextBools{"user_has_paid_bing_access": true})
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
// Test: feature flag true, but no user access
ctx2 := context.Background()
ctx2 = ContextWithFeatureChecker(ctx2, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx2 = ContextWithBools(ctx2, ContextBools{"user_has_paid_bing_access": false})
result2, err2 := cond.Evaluate(ctx2)
require.NoError(t, err2)
assert.False(t, result2)
// Test: feature flag false
ctx3 := context.Background()
ctx3 = ContextWithFeatureChecker(ctx3, func(_ context.Context, _ string) (bool, error) {
return false, nil
})
ctx3 = ContextWithBools(ctx3, ContextBools{"user_has_paid_bing_access": true})
result3, err3 := cond.Evaluate(ctx3)
require.NoError(t, err3)
assert.False(t, result3)
})
t.Run("CCA bypass pattern (CCA OR feature flag)", func(t *testing.T) {
// Pattern: CCA requests bypass feature flag, non-CCA requires feature flag
cond := Or(
ContextBool("is_cca"),
FeatureFlag("agent_search"),
)
// Test: CCA request (bypass feature flag)
ctx := context.Background()
ctx = ContextWithBools(ctx, ContextBools{"is_cca": true})
// No feature checker - CCA should pass without it
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
// Test: non-CCA with feature flag enabled
ctx2 := context.Background()
ctx2 = ContextWithFeatureChecker(ctx2, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx2 = ContextWithBools(ctx2, ContextBools{"is_cca": false})
result2, err2 := cond.Evaluate(ctx2)
require.NoError(t, err2)
assert.True(t, result2)
// Test: non-CCA with feature flag disabled
ctx3 := context.Background()
ctx3 = ContextWithFeatureChecker(ctx3, func(_ context.Context, _ string) (bool, error) {
return false, nil
})
ctx3 = ContextWithBools(ctx3, ContextBools{"is_cca": false})
result3, err3 := cond.Evaluate(ctx3)
require.NoError(t, err3)
assert.False(t, result3)
})
t.Run("CCA AND feature flag pattern", func(t *testing.T) {
// Pattern: must be CCA AND have feature flag enabled
cond := And(
ContextBool("is_cca"),
FeatureFlag("complex_workflows"),
)
// Test: CCA with feature flag
ctx := context.Background()
ctx = ContextWithFeatureChecker(ctx, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx = ContextWithBools(ctx, ContextBools{"is_cca": true})
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
// Test: CCA without feature flag
ctx2 := context.Background()
ctx2 = ContextWithFeatureChecker(ctx2, func(_ context.Context, _ string) (bool, error) {
return false, nil
})
ctx2 = ContextWithBools(ctx2, ContextBools{"is_cca": true})
result2, err2 := cond.Evaluate(ctx2)
require.NoError(t, err2)
assert.False(t, result2)
// Test: non-CCA with feature flag
ctx3 := context.Background()
ctx3 = ContextWithFeatureChecker(ctx3, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx3 = ContextWithBools(ctx3, ContextBools{"is_cca": false})
result3, err3 := cond.Evaluate(ctx3)
require.NoError(t, err3)
assert.False(t, result3)
})
t.Run("copilot-chat bypass pattern", func(t *testing.T) {
// Pattern: copilot-chat bypasses feature flag check
cond := Or(
ContextBool("mcp_host_is_copilot_chat"),
FeatureFlag("semantic_code_search"),
)
// Test: copilot-chat host (bypass feature flag)
ctx := context.Background()
ctx = ContextWithBools(ctx, ContextBools{"mcp_host_is_copilot_chat": true})
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
// Test: other host with feature flag
ctx2 := context.Background()
ctx2 = ContextWithFeatureChecker(ctx2, func(_ context.Context, _ string) (bool, error) {
return true, nil
})
ctx2 = ContextWithBools(ctx2, ContextBools{"mcp_host_is_copilot_chat": false})
result2, err2 := cond.Evaluate(ctx2)
require.NoError(t, err2)
assert.True(t, result2)
// Test: other host without feature flag
ctx3 := context.Background()
ctx3 = ContextWithFeatureChecker(ctx3, func(_ context.Context, _ string) (bool, error) {
return false, nil
})
ctx3 = ContextWithBools(ctx3, ContextBools{"mcp_host_is_copilot_chat": false})
result3, err3 := cond.Evaluate(ctx3)
require.NoError(t, err3)
assert.False(t, result3)
})
}
func TestContextBoolsMerging(t *testing.T) {
ctx := context.Background()
// Add first set of bools
ctx = ContextWithBools(ctx, ContextBools{"key1": true, "key2": false})
// Add second set - should merge
ctx = ContextWithBools(ctx, ContextBools{"key3": true, "key2": true}) // key2 overwritten
// Check all keys
assert.True(t, ContextBoolFromContext(ctx, "key1"))
assert.True(t, ContextBoolFromContext(ctx, "key2")) // overwritten value
assert.True(t, ContextBoolFromContext(ctx, "key3"))
assert.False(t, ContextBoolFromContext(ctx, "nonexistent"))
}
func TestSetContextBool(t *testing.T) {
ctx := context.Background()
ctx = SetContextBool(ctx, "my_flag", true)
assert.True(t, ContextBoolFromContext(ctx, "my_flag"))
assert.False(t, ContextBoolFromContext(ctx, "other_flag"))
}
func TestAndShortCircuit(t *testing.T) {
callCount := 0
ctx := context.Background()
// First condition returns false, second should not be called
cond := And(
Never(),
ConditionFunc(func(_ context.Context) (bool, error) {
callCount++
return true, nil
}),
)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.False(t, result)
assert.Equal(t, 0, callCount, "second condition should not be called due to short-circuit")
}
func TestOrShortCircuit(t *testing.T) {
callCount := 0
ctx := context.Background()
// First condition returns true, second should not be called
cond := Or(
Always(),
ConditionFunc(func(_ context.Context) (bool, error) {
callCount++
return false, nil
}),
)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
assert.Equal(t, 0, callCount, "second condition should not be called due to short-circuit")
}
func TestOrContinuesOnError(t *testing.T) {
ctx := context.Background()
// First condition errors, but second is true - should return true
cond := Or(
ConditionFunc(func(_ context.Context) (bool, error) {
return false, errors.New("error")
}),
Always(),
)
result, err := cond.Evaluate(ctx)
require.NoError(t, err)
assert.True(t, result)
}
File diff suppressed because it is too large Load Diff
+104 -42
View File
@@ -50,15 +50,61 @@ func (r *Inventory) isFeatureFlagAllowed(ctx context.Context, enableFlag, disabl
return true
}
// buildRequestMask creates a RequestMask for the current request context.
// This computes all condition values once for O(1) evaluation of each tool.
func (r *Inventory) buildRequestMask(ctx context.Context) *RequestMask {
if r.conditionCompiler == nil {
return nil
}
var bits uint64
bools := contextBoolsFromContext(ctx)
checker := FeatureCheckerFromContext(ctx)
r.conditionCompiler.mu.RLock()
defer r.conditionCompiler.mu.RUnlock()
for key, bit := range r.conditionCompiler.keyToBit {
// Keys are formatted as "ctx:key_name" or "ff:flag_name"
if len(key) < 4 { // Minimum: "ff:x" or "ctx:" prefix + 1 char
continue
}
switch {
case len(key) > 4 && key[:4] == "ctx:":
// Context bool: "ctx:key_name"
name := key[4:]
if bools != nil && bools[name] {
bits |= 1 << bit
}
case len(key) > 3 && key[:3] == "ff:":
// Feature flag: "ff:flag_name"
name := key[3:]
if checker != nil {
enabled, err := checker(ctx, name)
if err == nil && enabled {
bits |= 1 << bit
}
}
}
}
return &RequestMask{
bits: bits,
ctx: ctx,
}
}
// isToolEnabled checks if a specific tool is enabled based on current filters.
// Filter evaluation order:
// 1. Tool.Enabled (tool self-filtering)
// 2. FeatureFlagEnable/FeatureFlagDisable
// 3. Read-only filter
// 4. Builder filters (via WithFilter)
// 5. Toolset/additional tools
func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
// 1. Check tool's own Enabled function first
// 1. Tool.Enabled (legacy tool self-filtering - deprecated)
// 2. Tool.EnableCondition via compiled bitmask (O(1) evaluation)
// 3. FeatureFlagEnable/FeatureFlagDisable (legacy - deprecated)
// 4. Read-only filter
// 5. Builder filters (via WithFilter)
// 6. Toolset/additional tools
func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool, toolIndex int, rm *RequestMask) bool {
// 1. Check tool's legacy Enabled function first (for backward compatibility)
if tool.Enabled != nil {
enabled, err := tool.Enabled(ctx)
if err != nil {
@@ -69,15 +115,48 @@ func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
return false
}
}
// 2. Check feature flags
// 2. Check tool's EnableCondition via compiled bitmask (O(1) evaluation)
if toolIndex >= 0 && toolIndex < len(r.compiledConditions) && r.compiledConditions[toolIndex] != nil {
if rm != nil {
enabled, err := r.compiledConditions[toolIndex].Evaluate(rm)
if err != nil {
fmt.Fprintf(os.Stderr, "Tool.EnableCondition check error for %q: %v\n", tool.Tool.Name, err)
return false
}
if !enabled {
return false
}
} else if tool.EnableCondition != nil {
// Fallback to tree-based evaluation if no request mask
enabled, err := tool.EnableCondition.Evaluate(ctx)
if err != nil {
fmt.Fprintf(os.Stderr, "Tool.EnableCondition check error for %q: %v\n", tool.Tool.Name, err)
return false
}
if !enabled {
return false
}
}
} else if tool.EnableCondition != nil {
// Fallback to tree-based evaluation if no compiled condition
enabled, err := tool.EnableCondition.Evaluate(ctx)
if err != nil {
fmt.Fprintf(os.Stderr, "Tool.EnableCondition check error for %q: %v\n", tool.Tool.Name, err)
return false
}
if !enabled {
return false
}
}
// 3. Check legacy feature flags (for backward compatibility)
if !r.isFeatureFlagAllowed(ctx, tool.FeatureFlagEnable, tool.FeatureFlagDisable) {
return false
}
// 3. Check read-only filter (applies to all tools)
// 4. Check read-only filter (applies to all tools)
if r.readOnly && !tool.IsReadOnly() {
return false
}
// 4. Apply builder filters
// 5. Apply builder filters
for _, filter := range r.filters {
allowed, err := filter(ctx, tool)
if err != nil {
@@ -88,11 +167,11 @@ func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
return false
}
}
// 5. Check if tool is in additionalTools (bypasses toolset filter)
// 6. Check if tool is in additionalTools (bypasses toolset filter)
if r.additionalTools != nil && r.additionalTools[tool.Tool.Name] {
return true
}
// 5. Check toolset filter
// 6. Check toolset filter
if !r.isToolsetEnabled(tool.Toolset.ID) {
return false
}
@@ -102,30 +181,30 @@ func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
// AvailableTools returns the tools that pass all current filters,
// sorted deterministically by toolset ID, then tool name.
// The context is used for feature flag evaluation.
// Uses O(1) bitmask evaluation for EnableConditions when possible.
// Note: Tools are pre-sorted at build time, so filtering preserves order.
func (r *Inventory) AvailableTools(ctx context.Context) []ServerTool {
// Build request mask once for O(1) condition evaluation
rm := r.buildRequestMask(ctx)
// Tools are pre-sorted at build time; filtering preserves order
var result []ServerTool
for i := range r.tools {
tool := &r.tools[i]
if r.isToolEnabled(ctx, tool) {
if r.isToolEnabled(ctx, tool, i, rm) {
result = append(result, *tool)
}
}
// Sort deterministically: by toolset ID, then by tool name
sort.Slice(result, func(i, j int) bool {
if result[i].Toolset.ID != result[j].Toolset.ID {
return result[i].Toolset.ID < result[j].Toolset.ID
}
return result[i].Tool.Name < result[j].Tool.Name
})
return result
}
// AvailableResourceTemplates returns resource templates that pass all current filters,
// sorted deterministically by toolset ID, then template name.
// The context is used for feature flag evaluation.
// Note: Resources are pre-sorted at build time, so filtering preserves order.
func (r *Inventory) AvailableResourceTemplates(ctx context.Context) []ServerResourceTemplate {
// Resources are pre-sorted at build time; filtering preserves order
var result []ServerResourceTemplate
for i := range r.resourceTemplates {
res := &r.resourceTemplates[i]
@@ -138,21 +217,15 @@ func (r *Inventory) AvailableResourceTemplates(ctx context.Context) []ServerReso
}
}
// Sort deterministically: by toolset ID, then by template name
sort.Slice(result, func(i, j int) bool {
if result[i].Toolset.ID != result[j].Toolset.ID {
return result[i].Toolset.ID < result[j].Toolset.ID
}
return result[i].Template.Name < result[j].Template.Name
})
return result
}
// AvailablePrompts returns prompts that pass all current filters,
// sorted deterministically by toolset ID, then prompt name.
// The context is used for feature flag evaluation.
// Note: Prompts are pre-sorted at build time, so filtering preserves order.
func (r *Inventory) AvailablePrompts(ctx context.Context) []ServerPrompt {
// Prompts are pre-sorted at build time; filtering preserves order
var result []ServerPrompt
for i := range r.prompts {
prompt := &r.prompts[i]
@@ -165,14 +238,6 @@ func (r *Inventory) AvailablePrompts(ctx context.Context) []ServerPrompt {
}
}
// Sort deterministically: by toolset ID, then by prompt name
sort.Slice(result, func(i, j int) bool {
if result[i].Toolset.ID != result[j].Toolset.ID {
return result[i].Toolset.ID < result[j].Toolset.ID
}
return result[i].Prompt.Name < result[j].Prompt.Name
})
return result
}
@@ -221,7 +286,9 @@ func (r *Inventory) filterPromptsByName(name string) []ServerPrompt {
// ToolsForToolset returns all tools belonging to a specific toolset.
// This method bypasses the toolset enabled filter (for dynamic toolset registration),
// but still respects the read-only filter.
// Note: Tools are pre-sorted at build time, so filtering preserves order.
func (r *Inventory) ToolsForToolset(toolsetID ToolsetID) []ServerTool {
// Tools are pre-sorted at build time; filtering preserves order
var result []ServerTool
for i := range r.tools {
tool := &r.tools[i]
@@ -234,11 +301,6 @@ func (r *Inventory) ToolsForToolset(toolsetID ToolsetID) []ServerTool {
}
}
// Sort by tool name for deterministic order
sort.Slice(result, func(i, j int) bool {
return result[i].Tool.Name < result[j].Tool.Name
})
return result
}
+10 -1
View File
@@ -24,6 +24,7 @@ import (
// - Deterministic ordering for documentation generation
// - Lazy dependency injection during registration via RegisterAll()
// - Runtime toolset enabling for dynamic toolsets mode
// - O(1) EnableCondition evaluation via pre-compiled bitmasks
type Inventory struct {
// tools holds all tools in this group (ordered for iteration)
tools []ServerTool
@@ -40,6 +41,12 @@ type Inventory struct {
defaultToolsetIDs []ToolsetID // sorted list of default toolset IDs
toolsetDescriptions map[ToolsetID]string // toolset ID -> description
// Compiled conditions for O(1) EnableCondition evaluation (set during Build)
// Maps tool index → compiled condition (nil means always enabled)
compiledConditions []*CompiledCondition
// conditionCompiler used to compile conditions (shared across inventory)
conditionCompiler *ConditionCompiler
// Filters - these control what's returned by Available* methods
// readOnly when true filters out write tools
readOnly bool
@@ -110,7 +117,9 @@ func (r *Inventory) ForMCPRequest(method string, itemName string) *Inventory {
enabledToolsets: r.enabledToolsets, // shared, not modified
additionalTools: r.additionalTools, // shared, not modified
featureChecker: r.featureChecker,
filters: r.filters, // shared, not modified
filters: r.filters, // shared, not modified
compiledConditions: r.compiledConditions, // shared, not modified
conditionCompiler: r.conditionCompiler, // shared, not modified
unrecognizedToolsets: r.unrecognizedToolsets,
}
+358 -4
View File
@@ -7,6 +7,7 @@ import (
"testing"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/stretchr/testify/assert"
)
// testToolsetMetadata returns a ToolsetMetadata for testing
@@ -1598,10 +1599,11 @@ func TestFilteredToolsMatchesAvailableTools(t *testing.T) {
func TestFilteringOrder(t *testing.T) {
// Test that filters are applied in the correct order:
// 1. Tool.Enabled
// 2. Feature flags
// 3. Read-only
// 4. Builder filters
// 5. Toolset/additional tools
// 2. EnableCondition
// 3. Feature flags
// 4. Read-only
// 5. Builder filters
// 6. Toolset/additional tools
callOrder := []string{}
@@ -1643,3 +1645,355 @@ func TestFilteringOrder(t *testing.T) {
}
}
}
// Tests for EnableCondition integration
func TestEnableConditionSimple(t *testing.T) {
// Tool with EnableCondition that returns true
tool := mockTool("test_tool", "toolset1", true)
tool.EnableCondition = Always()
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
available := reg.AvailableTools(context.Background())
if len(available) != 1 {
t.Error("Tool should be included when EnableCondition returns true")
}
// Tool with EnableCondition that returns false
tool2 := mockTool("test_tool2", "toolset1", true)
tool2.EnableCondition = Never()
reg2 := NewBuilder().
SetTools([]ServerTool{tool2}).
WithToolsets([]string{"all"}).
Build()
available2 := reg2.AvailableTools(context.Background())
if len(available2) != 0 {
t.Error("Tool should be excluded when EnableCondition returns false")
}
}
func TestEnableConditionWithFeatureFlag(t *testing.T) {
// Tool with EnableCondition using FeatureFlag condition
tool := mockTool("test_tool", "toolset1", true)
tool.EnableCondition = FeatureFlag("my_feature")
// Without feature checker - should be excluded
reg1 := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
available1 := reg1.AvailableTools(context.Background())
if len(available1) != 0 {
t.Error("Tool should be excluded when no feature checker is available")
}
// With feature checker that returns true
ctx := ContextWithFeatureChecker(context.Background(), func(_ context.Context, flag string) (bool, error) {
return flag == "my_feature", nil
})
reg2 := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
available2 := reg2.AvailableTools(ctx)
if len(available2) != 1 {
t.Error("Tool should be included when feature flag is enabled via context")
}
}
func TestEnableConditionWithContextBool(t *testing.T) {
// Tool with EnableCondition using ContextBool
tool := mockTool("cca_tool", "toolset1", true)
tool.EnableCondition = ContextBool("is_cca")
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
// Without context bool - should be excluded
available1 := reg.AvailableTools(context.Background())
if len(available1) != 0 {
t.Error("Tool should be excluded when context bool is not set")
}
// With context bool = true
ctx := ContextWithBools(context.Background(), ContextBools{"is_cca": true})
available2 := reg.AvailableTools(ctx)
if len(available2) != 1 {
t.Error("Tool should be included when context bool is true")
}
// With context bool = false
ctx3 := ContextWithBools(context.Background(), ContextBools{"is_cca": false})
available3 := reg.AvailableTools(ctx3)
if len(available3) != 0 {
t.Error("Tool should be excluded when context bool is false")
}
}
func TestEnableConditionComplexPattern(t *testing.T) {
// CCA bypass pattern: tool is available if CCA OR feature flag is enabled
tool := mockTool("agent_tool", "toolset1", true)
tool.EnableCondition = Or(
ContextBool("is_cca"),
FeatureFlag("agent_search"),
)
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
// CCA request - should be enabled without feature flag
ctxCCA := ContextWithBools(context.Background(), ContextBools{"is_cca": true})
availableCCA := reg.AvailableTools(ctxCCA)
if len(availableCCA) != 1 {
t.Error("Tool should be enabled for CCA requests")
}
// Non-CCA with feature flag - should be enabled
ctxFF := ContextWithBools(context.Background(), ContextBools{"is_cca": false})
ctxFF = ContextWithFeatureChecker(ctxFF, func(_ context.Context, flag string) (bool, error) {
return flag == "agent_search", nil
})
availableFF := reg.AvailableTools(ctxFF)
if len(availableFF) != 1 {
t.Error("Tool should be enabled with feature flag for non-CCA")
}
// Non-CCA without feature flag - should be excluded
ctxNone := ContextWithBools(context.Background(), ContextBools{"is_cca": false})
availableNone := reg.AvailableTools(ctxNone)
if len(availableNone) != 0 {
t.Error("Tool should be excluded for non-CCA without feature flag")
}
}
func TestEnableConditionAndLegacyEnabledInteraction(t *testing.T) {
// When both Enabled and EnableCondition are set, both must pass
tool := mockTool("test_tool", "toolset1", true)
tool.Enabled = func(_ context.Context) (bool, error) {
return true, nil // Legacy Enabled passes
}
tool.EnableCondition = Never() // But EnableCondition fails
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
available := reg.AvailableTools(context.Background())
if len(available) != 0 {
t.Error("Tool should be excluded when EnableCondition returns false, even if Enabled returns true")
}
// Both pass
tool2 := mockTool("test_tool2", "toolset1", true)
tool2.Enabled = func(_ context.Context) (bool, error) {
return true, nil
}
tool2.EnableCondition = Always()
reg2 := NewBuilder().
SetTools([]ServerTool{tool2}).
WithToolsets([]string{"all"}).
Build()
available2 := reg2.AvailableTools(context.Background())
if len(available2) != 1 {
t.Error("Tool should be included when both Enabled and EnableCondition pass")
}
}
func TestEnableConditionFilteringOrder(t *testing.T) {
// Test that EnableCondition is checked after legacy Enabled but before feature flags
callOrder := []string{}
tool := mockToolWithFlags("test_tool", "toolset1", true, "my_feature", "")
tool.Enabled = func(_ context.Context) (bool, error) {
callOrder = append(callOrder, "LegacyEnabled")
return true, nil
}
tool.EnableCondition = ConditionFunc(func(_ context.Context) (bool, error) {
callOrder = append(callOrder, "EnableCondition")
return false, nil // Return false to stop early
})
checker := func(_ context.Context, _ string) (bool, error) {
callOrder = append(callOrder, "FeatureFlag")
return true, nil
}
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
WithFeatureChecker(checker).
Build()
_ = reg.AvailableTools(context.Background())
// Should stop at EnableCondition since it returns false
expectedOrder := []string{"LegacyEnabled", "EnableCondition"}
if len(callOrder) != len(expectedOrder) {
t.Errorf("Expected %d checks, got %d: %v", len(expectedOrder), len(callOrder), callOrder)
}
for i, expected := range expectedOrder {
if i >= len(callOrder) || callOrder[i] != expected {
t.Errorf("At position %d: expected %s, got %v", i, expected, callOrder)
}
}
}
func TestCompiledConditionsIntegration(t *testing.T) {
// Test that conditions are compiled at build time and evaluated with bitmask
// at request time (via AvailableTools)
checker := func(_ context.Context, flagName string) (bool, error) {
switch flagName {
case "enabled_flag":
return true, nil
case "disabled_flag":
return false, nil
default:
return false, nil
}
}
toolAlwaysEnabled := mockTool("always_on", "test", true)
toolAlwaysEnabled.EnableCondition = Always()
toolNeverEnabled := mockTool("never_on", "test", true)
toolNeverEnabled.EnableCondition = Never()
toolFlagEnabled := mockTool("flag_on", "test", true)
toolFlagEnabled.EnableCondition = FeatureFlag("enabled_flag")
toolFlagDisabled := mockTool("flag_off", "test", true)
toolFlagDisabled.EnableCondition = FeatureFlag("disabled_flag")
toolContextBool := mockTool("ctx_bool", "test", true)
toolContextBool.EnableCondition = ContextBool("my_bool")
toolComplex := mockTool("complex", "test", true)
toolComplex.EnableCondition = Or(
ContextBool("my_bool"),
FeatureFlag("enabled_flag"),
)
reg := NewBuilder().
SetTools([]ServerTool{
toolAlwaysEnabled,
toolNeverEnabled,
toolFlagEnabled,
toolFlagDisabled,
toolContextBool,
toolComplex,
}).
WithToolsets([]string{"all"}).
Build()
// Verify compiler was created
assert.NotNil(t, reg.conditionCompiler, "Condition compiler should be created")
assert.NotNil(t, reg.compiledConditions, "Compiled conditions should be created")
// Test without context bools - only flag-based tools should pass
ctx := context.Background()
ctx = ContextWithFeatureChecker(ctx, checker)
tools := reg.AvailableTools(ctx)
toolNames := make([]string, len(tools))
for i, tool := range tools {
toolNames[i] = tool.Tool.Name
}
assert.Contains(t, toolNames, "always_on", "Always enabled tool should be available")
assert.Contains(t, toolNames, "flag_on", "Flag enabled tool should be available")
assert.Contains(t, toolNames, "complex", "Complex (OR flag) tool should be available")
assert.NotContains(t, toolNames, "never_on", "Never enabled tool should not be available")
assert.NotContains(t, toolNames, "flag_off", "Disabled flag tool should not be available")
assert.NotContains(t, toolNames, "ctx_bool", "Context bool tool should not be available without bool set")
// Test with context bool set
ctx = ContextWithBools(ctx, ContextBools{"my_bool": true})
tools = reg.AvailableTools(ctx)
toolNames = make([]string, len(tools))
for i, tool := range tools {
toolNames[i] = tool.Tool.Name
}
assert.Contains(t, toolNames, "ctx_bool", "Context bool tool should be available when bool is set")
assert.Contains(t, toolNames, "complex", "Complex tool should still be available")
}
func TestCompiledConditionsANDBitmask(t *testing.T) {
// Test that AND conditions are compiled to bitmask AND operations
checker := func(_ context.Context, flagName string) (bool, error) {
return flagName == "flag_a" || flagName == "flag_b", nil
}
// Tool requires both flags
tool := mockTool("both_flags", "test", true)
tool.EnableCondition = And(
FeatureFlag("flag_a"),
FeatureFlag("flag_b"),
)
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
// Verify it was compiled to bitmask AND
assert.NotNil(t, reg.compiledConditions[0])
compiled := reg.compiledConditions[0]
assert.Equal(t, evalBitmaskAnd, compiled.evalType, "AND of flags should compile to bitmaskAnd")
// Test evaluation
ctx := context.Background()
ctx = ContextWithFeatureChecker(ctx, checker)
tools := reg.AvailableTools(ctx)
assert.Len(t, tools, 1, "Tool with both flags enabled should be available")
}
func TestCompiledConditionsORBitmask(t *testing.T) {
// Test that OR conditions are compiled to bitmask OR operations
checker := func(_ context.Context, flagName string) (bool, error) {
return flagName == "flag_a", nil // Only flag_a is enabled
}
// Tool requires either flag
tool := mockTool("either_flag", "test", true)
tool.EnableCondition = Or(
FeatureFlag("flag_a"),
FeatureFlag("flag_b"),
)
reg := NewBuilder().
SetTools([]ServerTool{tool}).
WithToolsets([]string{"all"}).
Build()
// Verify it was compiled to bitmask OR
assert.NotNil(t, reg.compiledConditions[0])
compiled := reg.compiledConditions[0]
assert.Equal(t, evalBitmaskOr, compiled.evalType, "OR of flags should compile to bitmaskOr")
// Test evaluation - should pass because flag_a is enabled
ctx := context.Background()
ctx = ContextWithFeatureChecker(ctx, checker)
tools := reg.AvailableTools(ctx)
assert.Len(t, tools, 1, "Tool with one flag enabled should be available")
}
+27
View File
@@ -58,17 +58,44 @@ type ServerTool struct {
// FeatureFlagEnable specifies a feature flag that must be enabled for this tool
// to be available. If set and the flag is not enabled, the tool is omitted.
//
// Deprecated: Use EnableCondition with FeatureFlag() instead for composable conditions.
// This field is checked before EnableCondition for backward compatibility.
FeatureFlagEnable string
// FeatureFlagDisable specifies a feature flag that, when enabled, causes this tool
// to be omitted. Used to disable tools when a feature flag is on.
//
// Deprecated: Use EnableCondition with Not(FeatureFlag()) instead for composable conditions.
// This field is checked before EnableCondition for backward compatibility.
FeatureFlagDisable string
// EnableCondition is the composable condition for tool availability.
// Use the condition combinators (FeatureFlag, ContextBool, And, Or, Not)
// to build complex enable logic declaratively.
//
// Examples:
// // Feature flag only
// EnableCondition: FeatureFlag("web_search")
//
// // Feature flag AND user policy
// EnableCondition: And(FeatureFlag("web_search"), ContextBool("user_has_paid_access"))
//
// // CCA bypass (CCA OR feature flag)
// EnableCondition: Or(ContextBool("is_cca"), FeatureFlag("agent_search"))
//
// If nil, the tool is enabled (subject to other filters like toolset, read-only).
EnableCondition EnableCondition
// Enabled is an optional function called at build/filter time to determine
// if this tool should be available. If nil, the tool is considered enabled
// (subject to FeatureFlagEnable/FeatureFlagDisable checks).
// The context carries request-scoped information for the consumer to use.
// Returns (enabled, error). On error, the tool should be treated as disabled.
//
// Deprecated: Use EnableCondition instead for composable, declarative conditions.
// If both Enabled and EnableCondition are set, Enabled takes precedence for
// backward compatibility. Migrate to EnableCondition for new tools.
Enabled func(ctx context.Context) (bool, error)
}