Skip to content

Commit 69ef0d4

Browse files
evalstateevalstate
authored
Merge pull request #113 from huggingface/claude/fix-tool-selection-logic-011CUzebrLKFe6QLeHRYGQmi
Review and fix tool selection logic
2 parents 1d0cfa8 + d3621b5 commit 69ef0d4

File tree

3 files changed

+163
-7
lines changed

3 files changed

+163
-7
lines changed

‎packages/app/src/server/mcp-proxy.ts‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -129,18 +129,19 @@ export const createProxyServerFactory = (
129129
const noImageFromSettings = settings?.builtInTools?.includes(GRADIO_IMAGE_FILTER_FLAG) ?? false;
130130
const stripImageContent = noImageFromHeader || noImageFromSettings;
131131

132-
// Skip Gradio endpoints if bouquet is not "all"
133-
if (bouquet && bouquet !== 'all') {
134-
logger.debug({ bouquet }, 'Bouquet specified and not "all", skipping Gradio endpoints');
135-
return result;
136-
}
137-
138132
// Skip Gradio endpoints if explicitly disabled
139133
if (gradio === 'none') {
140134
logger.debug('Gradio endpoints explicitly disabled via gradio="none"');
141135
return result;
142136
}
143137

138+
// Skip Gradio endpoints from settings if bouquet is specified (not "all") and no explicit gradio param
139+
// This allows: bouquet=search&gradio=foo/bar to work as expected
140+
if (bouquet && bouquet !== 'all' && !gradio) {
141+
logger.debug({ bouquet }, 'Bouquet specified (not "all") and no explicit gradio param, skipping Gradio endpoints from settings');
142+
return result;
143+
}
144+
144145
// Now we have access to userDetails if needed
145146
if (userDetails) {
146147
logger.debug(`Proxy has access to user details for: ${userDetails.name}`);

‎packages/app/src/server/utils/tool-selection-strategy.ts‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,40 @@ export interface ToolSelectionResult {
3232

3333
/**
3434
* Tool Selection Strategy - implements clear precedence rules for tool selection
35+
*
36+
* ## Two Independent Systems
37+
*
38+
* 1. **Built-in Tool Selection** (handled by this class):
39+
* - Controlled by: bouquet, mix, and user settings
40+
* - Affects: Built-in HuggingFace MCP tools only
41+
* - Returns: enabledToolIds array
42+
*
43+
* 2. **Gradio Endpoint Registration** (handled by mcp-proxy.ts):
44+
* - Controlled by: gradio parameter + user settings spaceTools
45+
* - Affects: Dynamic Gradio Space endpoints
46+
* - Works with any bouquet when explicitly specified via gradio=
47+
* - Special: gradio="none" disables all Gradio endpoints
48+
*
49+
* ## Precedence for Built-in Tools
50+
*
51+
* 1. BOUQUET_OVERRIDE (highest) - Completely replaces tool selection
52+
* 2. MIX - Adds mix bouquet tools to user settings
53+
* 3. USER_SETTINGS - Uses external or internal API settings
54+
* 4. FALLBACK (lowest) - All tools enabled when no config available
55+
*
56+
* ## Gradio Parameter Behavior
57+
*
58+
* - When `gradio=foo/bar` is **explicitly specified**, those endpoints are always included
59+
* - When `bouquet=search` (no gradio param), Gradio endpoints from settings are skipped
60+
* - When `bouquet=all`, Gradio endpoints from user settings are included
61+
* - Examples:
62+
* - `bouquet=search&gradio=microsoft/Florence-2-large` → search tools + Florence endpoint ✓
63+
* - `bouquet=hf_api&gradio=foo/bar` → hf_api tools + foo/bar endpoint ✓
64+
* - `bouquet=search` (no gradio param) → search tools only ✓
65+
* - `bouquet=all` → all tools + gradio endpoints from settings ✓
66+
*
67+
* The gradio parameter is parsed here for metadata/logging purposes only.
68+
* Actual endpoint registration happens in mcp-proxy.ts.
3569
*/
3670
export class ToolSelectionStrategy {
3771
private apiClient: McpApiClient;
@@ -87,11 +121,16 @@ export class ToolSelectionStrategy {
87121
* 2. Mix + user settings (additive)
88122
* 3. User settings (external/internal API)
89123
* 4. Fallback (all tools)
124+
*
125+
* Note: The `gradio` parameter is parsed and included in the result regardless of
126+
* the bouquet/mix/settings selection. The actual endpoint registration in mcp-proxy.ts
127+
* will respect the explicit gradio parameter even when a non-"all" bouquet is specified.
90128
*/
91129
async selectTools(context: ToolSelectionContext): Promise<ToolSelectionResult> {
92130
const { bouquet, mix, gradio } = extractAuthBouquetAndMix(context.headers);
93131

94-
// Parse gradio endpoints if provided
132+
// Parse gradio endpoints if provided (independent of bouquet selection)
133+
// These endpoints will be registered in mcp-proxy.ts unless gradio="none"
95134
const gradioSpaceTools = gradio ? this.parseGradioEndpoints(gradio) : [];
96135

97136
// 1. Bouquet override (highest precedence)

‎packages/app/test/server/utils/tool-selection-strategy.test.ts‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -681,4 +681,120 @@ describe('ToolSelectionStrategy', () => {
681681
expect(result.enabledToolIds).toContain('hf_doc_fetch');
682682
});
683683
});
684+
685+
describe('Gradio endpoint handling', () => {
686+
it('should include gradio endpoints in bouquet override mode', async () => {
687+
const context: ToolSelectionContext = {
688+
headers: {
689+
'x-mcp-bouquet': 'search',
690+
'x-mcp-gradio': 'microsoft/Florence-2-large,meta-llama/Llama-2-7b-chat-hf',
691+
},
692+
hfToken: 'test-token',
693+
};
694+
695+
const result = await strategy.selectTools(context);
696+
697+
expect(result.mode).toBe(ToolSelectionMode.BOUQUET_OVERRIDE);
698+
expect(result.enabledToolIds).toEqual(TOOL_ID_GROUPS.search);
699+
expect(result.reason).toBe('Bouquet override: search + 2 gradio endpoints');
700+
expect(result.gradioSpaceTools).toBeDefined();
701+
expect(result.gradioSpaceTools).toHaveLength(2);
702+
expect(result.gradioSpaceTools?.[0].name).toBe('microsoft/Florence-2-large');
703+
expect(result.gradioSpaceTools?.[1].name).toBe('meta-llama/Llama-2-7b-chat-hf');
704+
});
705+
706+
it('should include gradio endpoints in mix mode', async () => {
707+
const userSettings: AppSettings = {
708+
builtInTools: ['hf_whoami'],
709+
spaceTools: [],
710+
};
711+
712+
const context: ToolSelectionContext = {
713+
headers: {
714+
'x-mcp-mix': 'hf_api',
715+
'x-mcp-gradio': 'foo/bar',
716+
},
717+
userSettings,
718+
hfToken: 'test-token',
719+
};
720+
721+
const result = await strategy.selectTools(context);
722+
723+
expect(result.mode).toBe(ToolSelectionMode.MIX);
724+
expect(result.reason).toBe('User settings + mix(hf_api) + 1 gradio endpoints');
725+
expect(result.gradioSpaceTools).toBeDefined();
726+
expect(result.gradioSpaceTools).toHaveLength(1);
727+
expect(result.gradioSpaceTools?.[0].name).toBe('foo/bar');
728+
});
729+
730+
it('should include gradio endpoints in user settings mode', async () => {
731+
const userSettings: AppSettings = {
732+
builtInTools: ['hf_semantic_search'],
733+
spaceTools: [],
734+
};
735+
736+
const context: ToolSelectionContext = {
737+
headers: {
738+
'x-mcp-gradio': 'test/space',
739+
},
740+
userSettings,
741+
hfToken: 'test-token',
742+
};
743+
744+
const result = await strategy.selectTools(context);
745+
746+
expect(result.mode).toBe(ToolSelectionMode.INTERNAL_API);
747+
expect(result.reason).toBe('Internal API user settings + 1 gradio endpoints');
748+
expect(result.gradioSpaceTools).toBeDefined();
749+
expect(result.gradioSpaceTools).toHaveLength(1);
750+
});
751+
752+
it('should include gradio endpoints in fallback mode', async () => {
753+
const context: ToolSelectionContext = {
754+
headers: {
755+
'x-mcp-gradio': 'fallback/test',
756+
},
757+
hfToken: 'test-token',
758+
};
759+
760+
const result = await strategy.selectTools(context);
761+
762+
expect(result.mode).toBe(ToolSelectionMode.FALLBACK);
763+
expect(result.reason).toBe('Fallback - no settings available + 1 gradio endpoints');
764+
expect(result.gradioSpaceTools).toBeDefined();
765+
expect(result.gradioSpaceTools).toHaveLength(1);
766+
});
767+
768+
it('should not include gradio endpoints when not specified', async () => {
769+
const context: ToolSelectionContext = {
770+
headers: { 'x-mcp-bouquet': 'search' },
771+
hfToken: 'test-token',
772+
};
773+
774+
const result = await strategy.selectTools(context);
775+
776+
expect(result.mode).toBe(ToolSelectionMode.BOUQUET_OVERRIDE);
777+
expect(result.reason).toBe('Bouquet override: search');
778+
expect(result.gradioSpaceTools).toBeUndefined();
779+
});
780+
781+
it('should handle multiple gradio endpoints with various formats', async () => {
782+
const context: ToolSelectionContext = {
783+
headers: {
784+
'x-mcp-bouquet': 'hf_api',
785+
'x-mcp-gradio': 'user/space-one,org/space-two',
786+
},
787+
hfToken: 'test-token',
788+
};
789+
790+
const result = await strategy.selectTools(context);
791+
792+
expect(result.gradioSpaceTools).toBeDefined();
793+
expect(result.gradioSpaceTools).toHaveLength(2);
794+
expect(result.gradioSpaceTools?.map(s => s.name)).toEqual([
795+
'user/space-one',
796+
'org/space-two',
797+
]);
798+
});
799+
});
684800
});

0 commit comments

Comments
 (0)