Skip to content

Commit 3a52f7b

Browse files
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.
1 parent 3c453dd commit 3a52f7b

File tree

10 files changed

+3429
-51
lines changed

10 files changed

+3429
-51
lines changed

pkg/inventory/builder.go

Lines changed: 83 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -128,14 +128,22 @@ func (b *Builder) WithFilter(filter ToolFilter) *Builder {
128128
}
129129

130130
// Build creates the final Inventory with all configuration applied.
131-
// This processes toolset filtering, tool name resolution, and sets up
131+
// This processes toolset filtering, tool name resolution, compiles EnableConditions
132+
// for O(1) evaluation, pre-sorts all items for deterministic output, and sets up
132133
// the inventory for use. The returned Inventory is ready for use with
133134
// AvailableTools(), RegisterAll(), etc.
134135
func (b *Builder) Build() *Inventory {
136+
// Pre-sort tools, resources, and prompts at build time.
137+
// This eliminates sorting overhead on every Available*() call.
138+
// Filtering preserves order, so if input is sorted, output is sorted.
139+
sortedTools := b.preSortTools()
140+
sortedResources := b.preSortResources()
141+
sortedPrompts := b.preSortPrompts()
142+
135143
r := &Inventory{
136-
tools: b.tools,
137-
resourceTemplates: b.resourceTemplates,
138-
prompts: b.prompts,
144+
tools: sortedTools,
145+
resourceTemplates: sortedResources,
146+
prompts: sortedPrompts,
139147
deprecatedAliases: b.deprecatedAliases,
140148
readOnly: b.readOnly,
141149
featureChecker: b.featureChecker,
@@ -158,9 +166,80 @@ func (b *Builder) Build() *Inventory {
158166
}
159167
}
160168

169+
// Compile EnableConditions for O(1) bitmask evaluation
170+
// Note: compileConditions uses r.tools which is now sortedTools
171+
r.conditionCompiler, r.compiledConditions = b.compileConditions(sortedTools)
172+
161173
return r
162174
}
163175

176+
// preSortTools returns a copy of tools sorted by toolset ID, then tool name.
177+
// This allows filtering to preserve order without re-sorting.
178+
func (b *Builder) preSortTools() []ServerTool {
179+
if len(b.tools) == 0 {
180+
return b.tools
181+
}
182+
sorted := make([]ServerTool, len(b.tools))
183+
copy(sorted, b.tools)
184+
sort.Slice(sorted, func(i, j int) bool {
185+
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
186+
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
187+
}
188+
return sorted[i].Tool.Name < sorted[j].Tool.Name
189+
})
190+
return sorted
191+
}
192+
193+
// preSortResources returns a copy of resources sorted by toolset ID, then template name.
194+
func (b *Builder) preSortResources() []ServerResourceTemplate {
195+
if len(b.resourceTemplates) == 0 {
196+
return b.resourceTemplates
197+
}
198+
sorted := make([]ServerResourceTemplate, len(b.resourceTemplates))
199+
copy(sorted, b.resourceTemplates)
200+
sort.Slice(sorted, func(i, j int) bool {
201+
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
202+
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
203+
}
204+
return sorted[i].Template.Name < sorted[j].Template.Name
205+
})
206+
return sorted
207+
}
208+
209+
// preSortPrompts returns a copy of prompts sorted by toolset ID, then prompt name.
210+
func (b *Builder) preSortPrompts() []ServerPrompt {
211+
if len(b.prompts) == 0 {
212+
return b.prompts
213+
}
214+
sorted := make([]ServerPrompt, len(b.prompts))
215+
copy(sorted, b.prompts)
216+
sort.Slice(sorted, func(i, j int) bool {
217+
if sorted[i].Toolset.ID != sorted[j].Toolset.ID {
218+
return sorted[i].Toolset.ID < sorted[j].Toolset.ID
219+
}
220+
return sorted[i].Prompt.Name < sorted[j].Prompt.Name
221+
})
222+
return sorted
223+
}
224+
225+
// compileConditions compiles all EnableConditions into bitmask-based evaluators.
226+
// Returns the compiler (for building request masks) and compiled conditions slice.
227+
// Takes the sorted tools slice to ensure compiled conditions align with sorted order.
228+
func (b *Builder) compileConditions(sortedTools []ServerTool) (*ConditionCompiler, []*CompiledCondition) {
229+
compiler := NewConditionCompiler()
230+
compiled := make([]*CompiledCondition, len(sortedTools))
231+
232+
for i := range sortedTools {
233+
if sortedTools[i].EnableCondition != nil {
234+
compiled[i] = compiler.Compile(sortedTools[i].EnableCondition)
235+
}
236+
// nil means no condition (always enabled from condition perspective)
237+
}
238+
239+
compiler.Freeze()
240+
return compiler, compiled
241+
}
242+
164243
// processToolsets processes the toolsetIDs configuration and returns:
165244
// - enabledToolsets map (nil means all enabled)
166245
// - unrecognizedToolsets list for warnings

0 commit comments

Comments
 (0)