Skip to content

Commit 6c1bde3

Browse files
Add comprehensive tests for feature-flagged tool handling
Co-authored-by: SamMorrowDrums <4811358+SamMorrowDrums@users.noreply.github.com>
1 parent 43bc950 commit 6c1bde3

File tree

2 files changed

+385
-0
lines changed

2 files changed

+385
-0
lines changed

pkg/github/tools_validation_test.go

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,3 +184,141 @@ func TestToolsetMetadataConsistency(t *testing.T) {
184184
}
185185
}
186186
}
187+
188+
// TestFeatureFlaggedToolsAreMutuallyExclusive validates that when multiple tools share
189+
// the same name with different feature flags, they are mutually exclusive (one uses
190+
// FeatureFlagEnable while the other uses FeatureFlagDisable with the same flag name).
191+
// This ensures there are no conflicts where two tools with the same name could be active
192+
// simultaneously, which would cause "unknown tool" errors.
193+
func TestFeatureFlaggedToolsAreMutuallyExclusive(t *testing.T) {
194+
tools := AllTools(stubTranslation)
195+
196+
// Group tools by name
197+
toolsByName := make(map[string][]inventory.ServerTool)
198+
for _, tool := range tools {
199+
toolsByName[tool.Tool.Name] = append(toolsByName[tool.Tool.Name], tool)
200+
}
201+
202+
// Check each group of tools with the same name
203+
for name, group := range toolsByName {
204+
if len(group) <= 1 {
205+
continue // Single tool, no conflict possible
206+
}
207+
208+
// Skip explicitly allowed duplicates (like get_label)
209+
if name == "get_label" {
210+
continue
211+
}
212+
213+
t.Run(name, func(t *testing.T) {
214+
// All tools with the same name must have feature flags
215+
for i, tool := range group {
216+
hasFlags := tool.FeatureFlagEnable != "" || tool.FeatureFlagDisable != ""
217+
assert.True(t, hasFlags,
218+
"Tool %q (variant %d/%d) shares a name with other tools but has no feature flags",
219+
name, i+1, len(group))
220+
}
221+
222+
// Verify mutual exclusivity: collect all flags used
223+
enableFlags := make(map[string]bool)
224+
disableFlags := make(map[string]bool)
225+
226+
for _, tool := range group {
227+
if tool.FeatureFlagEnable != "" {
228+
enableFlags[tool.FeatureFlagEnable] = true
229+
}
230+
if tool.FeatureFlagDisable != "" {
231+
disableFlags[tool.FeatureFlagDisable] = true
232+
}
233+
}
234+
235+
// For proper mutual exclusivity, we expect:
236+
// - Same flag name used in both FeatureFlagEnable and FeatureFlagDisable
237+
// - This ensures only one variant is active at a time
238+
if len(group) == 2 {
239+
// Most common case: two variants controlled by one flag
240+
var flagName string
241+
for flag := range enableFlags {
242+
flagName = flag
243+
break
244+
}
245+
for flag := range disableFlags {
246+
if flagName == "" {
247+
flagName = flag
248+
}
249+
}
250+
251+
if flagName != "" {
252+
assert.True(t, enableFlags[flagName] || disableFlags[flagName],
253+
"Tools sharing name %q should use the same flag for mutual exclusivity", name)
254+
255+
// Verify exactly one tool has FeatureFlagEnable=flagName and one has FeatureFlagDisable=flagName
256+
enableCount := 0
257+
disableCount := 0
258+
for _, tool := range group {
259+
if tool.FeatureFlagEnable == flagName {
260+
enableCount++
261+
}
262+
if tool.FeatureFlagDisable == flagName {
263+
disableCount++
264+
}
265+
}
266+
267+
// We should have at least one of each for proper mutual exclusivity
268+
assert.True(t, enableCount > 0 || disableCount > 0,
269+
"Tools sharing name %q should have proper feature flag configuration", name)
270+
}
271+
}
272+
273+
// Additional safety check: ensure no two tools could be enabled simultaneously
274+
// by testing all possible flag states
275+
testFlagStates := []map[string]bool{
276+
{}, // All flags off
277+
// We'll add specific flag combinations based on what flags we found
278+
}
279+
280+
// Add test cases for each unique flag
281+
allFlags := make(map[string]bool)
282+
for flag := range enableFlags {
283+
allFlags[flag] = true
284+
}
285+
for flag := range disableFlags {
286+
allFlags[flag] = true
287+
}
288+
289+
for flag := range allFlags {
290+
testFlagStates = append(testFlagStates, map[string]bool{flag: true})
291+
}
292+
293+
// Test each flag state
294+
for stateIdx, flagState := range testFlagStates {
295+
enabledCount := 0
296+
for _, tool := range group {
297+
enabled := true
298+
299+
// Check FeatureFlagEnable - tool requires this flag to be ON
300+
if tool.FeatureFlagEnable != "" {
301+
if !flagState[tool.FeatureFlagEnable] {
302+
enabled = false
303+
}
304+
}
305+
306+
// Check FeatureFlagDisable - tool is disabled if this flag is ON
307+
if tool.FeatureFlagDisable != "" {
308+
if flagState[tool.FeatureFlagDisable] {
309+
enabled = false
310+
}
311+
}
312+
313+
if enabled {
314+
enabledCount++
315+
}
316+
}
317+
318+
assert.LessOrEqual(t, enabledCount, 1,
319+
"Flag state %d for tool %q: At most one variant should be enabled (found %d enabled)",
320+
stateIdx, name, enabledCount)
321+
}
322+
})
323+
}
324+
}

pkg/inventory/registry_test.go

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1688,3 +1688,250 @@ func TestForMCPRequest_ToolsCall_FeatureFlaggedVariants(t *testing.T) {
16881688
availableOn[0].FeatureFlagEnable, availableOn[0].FeatureFlagDisable)
16891689
}
16901690
}
1691+
1692+
// TestToolsList_WithFeatureFlags validates that tools/list returns only the tools
1693+
// available based on the current feature flag state, without duplicates
1694+
func TestToolsList_WithFeatureFlags(t *testing.T) {
1695+
// Create tools with various feature flag configurations
1696+
// These are properly mutually exclusive
1697+
tools := []ServerTool{
1698+
mockToolWithFlags("tool_a", "test", true, "", "flag_x"), // disabled when flag_x is ON
1699+
mockToolWithFlags("tool_a", "test", true, "flag_x", ""), // enabled when flag_x is ON
1700+
mockToolWithFlags("tool_b", "test", true, "flag_y", ""), // enabled only when flag_y is ON
1701+
mockToolWithFlags("tool_c", "test", true, "", "flag_z"), // disabled when flag_z is ON
1702+
mockToolWithFlags("tool_c", "test", true, "flag_z", ""), // enabled when flag_z is ON
1703+
mockTool("tool_d", "test", true), // always enabled (no flags)
1704+
}
1705+
1706+
testCases := []struct {
1707+
name string
1708+
flagStates map[string]bool
1709+
expectedTools []string // tool names that should be available
1710+
}{
1711+
{
1712+
name: "All flags OFF",
1713+
flagStates: map[string]bool{},
1714+
expectedTools: []string{"tool_a", "tool_c", "tool_d"},
1715+
},
1716+
{
1717+
name: "flag_x ON",
1718+
flagStates: map[string]bool{"flag_x": true},
1719+
expectedTools: []string{"tool_a", "tool_c", "tool_d"},
1720+
},
1721+
{
1722+
name: "flag_y ON",
1723+
flagStates: map[string]bool{"flag_y": true},
1724+
expectedTools: []string{"tool_a", "tool_b", "tool_c", "tool_d"},
1725+
},
1726+
{
1727+
name: "flag_z ON",
1728+
flagStates: map[string]bool{"flag_z": true},
1729+
expectedTools: []string{"tool_a", "tool_c", "tool_d"},
1730+
},
1731+
{
1732+
name: "flag_x and flag_y ON",
1733+
flagStates: map[string]bool{"flag_x": true, "flag_y": true},
1734+
expectedTools: []string{"tool_a", "tool_b", "tool_c", "tool_d"},
1735+
},
1736+
}
1737+
1738+
for _, tc := range testCases {
1739+
t.Run(tc.name, func(t *testing.T) {
1740+
// Create feature checker that returns the flag states for this test case
1741+
checker := func(_ context.Context, flag string) (bool, error) {
1742+
return tc.flagStates[flag], nil
1743+
}
1744+
1745+
reg := NewBuilder().
1746+
SetTools(tools).
1747+
WithToolsets([]string{"all"}).
1748+
WithFeatureChecker(checker).
1749+
Build()
1750+
1751+
// Test tools/list endpoint
1752+
listReg := reg.ForMCPRequest(MCPMethodToolsList, "")
1753+
available := listReg.AvailableTools(context.Background())
1754+
1755+
// Collect available tool names
1756+
availableNames := make(map[string]int)
1757+
for _, tool := range available {
1758+
availableNames[tool.Tool.Name]++
1759+
}
1760+
1761+
// Verify expected tools are present
1762+
for _, expectedName := range tc.expectedTools {
1763+
count, found := availableNames[expectedName]
1764+
if !found {
1765+
t.Errorf("Expected tool %q not found in available tools", expectedName)
1766+
} else if count > 1 {
1767+
t.Errorf("Tool %q appears %d times (should appear only once)", expectedName, count)
1768+
}
1769+
}
1770+
1771+
// Verify no unexpected tools
1772+
if len(availableNames) != len(tc.expectedTools) {
1773+
t.Errorf("Expected %d tools, got %d: %v", len(tc.expectedTools), len(availableNames), availableNames)
1774+
}
1775+
1776+
// Verify no duplicate tool names in the result
1777+
for name, count := range availableNames {
1778+
if count > 1 {
1779+
t.Errorf("Duplicate tool name %q appears %d times", name, count)
1780+
}
1781+
}
1782+
})
1783+
}
1784+
}
1785+
1786+
// TestToolsCall_WithFeatureFlags validates that tools/call (ForMCPRequest with specific tool)
1787+
// returns the correct tool variant based on feature flags
1788+
func TestToolsCall_WithFeatureFlags(t *testing.T) {
1789+
tools := []ServerTool{
1790+
mockToolWithFlags("shared_tool", "test", true, "", "feature_flag"), // OLD: disabled when feature_flag is ON
1791+
mockToolWithFlags("shared_tool", "test", true, "feature_flag", ""), // NEW: enabled when feature_flag is ON
1792+
mockTool("other_tool", "test", true),
1793+
}
1794+
1795+
testCases := []struct {
1796+
name string
1797+
toolName string
1798+
featureFlagOn bool
1799+
expectToolCount int
1800+
expectEnableFlag string
1801+
expectDisableFlag string
1802+
}{
1803+
{
1804+
name: "Call shared_tool with flag OFF - should get old variant",
1805+
toolName: "shared_tool",
1806+
featureFlagOn: false,
1807+
expectToolCount: 1,
1808+
expectEnableFlag: "",
1809+
expectDisableFlag: "feature_flag",
1810+
},
1811+
{
1812+
name: "Call shared_tool with flag ON - should get new variant",
1813+
toolName: "shared_tool",
1814+
featureFlagOn: true,
1815+
expectToolCount: 1,
1816+
expectEnableFlag: "feature_flag",
1817+
expectDisableFlag: "",
1818+
},
1819+
{
1820+
name: "Call other_tool - always available",
1821+
toolName: "other_tool",
1822+
featureFlagOn: false,
1823+
expectToolCount: 1,
1824+
expectEnableFlag: "",
1825+
expectDisableFlag: "",
1826+
},
1827+
}
1828+
1829+
for _, tc := range testCases {
1830+
t.Run(tc.name, func(t *testing.T) {
1831+
var checker FeatureFlagChecker
1832+
if tc.featureFlagOn {
1833+
checker = func(_ context.Context, flag string) (bool, error) {
1834+
return flag == "feature_flag", nil
1835+
}
1836+
} else {
1837+
checker = func(_ context.Context, _ string) (bool, error) {
1838+
return false, nil
1839+
}
1840+
}
1841+
1842+
reg := NewBuilder().
1843+
SetTools(tools).
1844+
WithToolsets([]string{"all"}).
1845+
WithFeatureChecker(checker).
1846+
Build()
1847+
1848+
// Test tools/call endpoint
1849+
callReg := reg.ForMCPRequest(MCPMethodToolsCall, tc.toolName)
1850+
available := callReg.AvailableTools(context.Background())
1851+
1852+
if len(available) != tc.expectToolCount {
1853+
t.Fatalf("Expected %d tool(s), got %d", tc.expectToolCount, len(available))
1854+
}
1855+
1856+
if tc.expectToolCount > 0 {
1857+
tool := available[0]
1858+
if tool.Tool.Name != tc.toolName {
1859+
t.Errorf("Expected tool name %q, got %q", tc.toolName, tool.Tool.Name)
1860+
}
1861+
if tool.FeatureFlagEnable != tc.expectEnableFlag {
1862+
t.Errorf("Expected FeatureFlagEnable=%q, got %q", tc.expectEnableFlag, tool.FeatureFlagEnable)
1863+
}
1864+
if tool.FeatureFlagDisable != tc.expectDisableFlag {
1865+
t.Errorf("Expected FeatureFlagDisable=%q, got %q", tc.expectDisableFlag, tool.FeatureFlagDisable)
1866+
}
1867+
}
1868+
})
1869+
}
1870+
}
1871+
1872+
// TestNoDuplicateToolsInAnyFeatureFlagCombination validates that no matter what
1873+
// combination of feature flags is enabled, we never have duplicate tool names in
1874+
// the available tools list
1875+
func TestNoDuplicateToolsInAnyFeatureFlagCombination(t *testing.T) {
1876+
tools := []ServerTool{
1877+
// Simulate real tools with feature flags
1878+
mockToolWithFlags("actions_list", "test", true, "", "consolidated"),
1879+
mockToolWithFlags("actions_list", "test", true, "consolidated", ""),
1880+
mockToolWithFlags("actions_get", "test", true, "", "consolidated"),
1881+
mockToolWithFlags("actions_get", "test", true, "consolidated", ""),
1882+
mockToolWithFlags("get_job_logs", "test", true, "", "consolidated"),
1883+
mockToolWithFlags("get_job_logs", "test", true, "consolidated", ""),
1884+
mockTool("regular_tool", "test", true),
1885+
mockToolWithFlags("feature_tool", "test", true, "other_flag", ""),
1886+
}
1887+
1888+
// Test all combinations of feature flags
1889+
flags := []string{"consolidated", "other_flag"}
1890+
1891+
// Generate all possible combinations of flags (2^n combinations)
1892+
numCombinations := 1 << len(flags)
1893+
1894+
for i := 0; i < numCombinations; i++ {
1895+
flagStates := make(map[string]bool)
1896+
var testName string
1897+
for j, flag := range flags {
1898+
isOn := (i & (1 << j)) != 0
1899+
flagStates[flag] = isOn
1900+
if isOn {
1901+
if testName != "" {
1902+
testName += "_"
1903+
}
1904+
testName += flag
1905+
}
1906+
}
1907+
if testName == "" {
1908+
testName = "no_flags"
1909+
}
1910+
1911+
t.Run(testName, func(t *testing.T) {
1912+
checker := func(_ context.Context, flag string) (bool, error) {
1913+
return flagStates[flag], nil
1914+
}
1915+
1916+
reg := NewBuilder().
1917+
SetTools(tools).
1918+
WithToolsets([]string{"all"}).
1919+
WithFeatureChecker(checker).
1920+
Build()
1921+
1922+
available := reg.AvailableTools(context.Background())
1923+
1924+
// Check for duplicates
1925+
seen := make(map[string]int)
1926+
for _, tool := range available {
1927+
seen[tool.Tool.Name]++
1928+
}
1929+
1930+
for name, count := range seen {
1931+
if count > 1 {
1932+
t.Errorf("Duplicate tool %q appears %d times with flag state: %v", name, count, flagStates)
1933+
}
1934+
}
1935+
})
1936+
}
1937+
}

0 commit comments

Comments
 (0)