hf-mcp-server/packages/app/test/server/utils/gradio-utils.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import * as hub from '@huggingface/hub';
import {
	isGradioTool,
	createGradioToolName,
	parseGradioSpaceIds,
	resolveMcpFileListingSource,
} from '../../../src/server/utils/gradio-utils.js';

vi.mock('@huggingface/hub', async (importOriginal) => {
	const actual = await importOriginal<typeof import('@huggingface/hub')>();
	return {
		...actual,
		repoExists: vi.fn(),
	};
});

describe('isGradioTool', () => {
	describe('should return true for valid Gradio tool names', () => {
		it('should detect basic gr tools', () => {
			expect(isGradioTool('gr1_tool')).toBe(true);
			expect(isGradioTool('gr2_another_tool')).toBe(true);
			expect(isGradioTool('gr999_test')).toBe(true);
		});

		it('should detect grp (private) tools', () => {
			expect(isGradioTool('grp1_private_tool')).toBe(true);
			expect(isGradioTool('grp2_another_private')).toBe(true);
			expect(isGradioTool('grp10_test')).toBe(true);
		});

		it('should detect real-world Gradio tool names', () => {
			expect(isGradioTool('gr1_evalstate_flux1_schnell')).toBe(true);
			expect(isGradioTool('grp3_my_private_space')).toBe(true);
			expect(isGradioTool('gr42_image_generator')).toBe(true);
		});

		it('should detect tool names generated by createGradioToolName', () => {
			// Generate examples using the actual tool creation logic
			const publicTool1 = createGradioToolName('flux1_schnell', 0, false);
			const publicTool2 = createGradioToolName('EasyGhibli', 1, false);
			const privateTool = createGradioToolName('private-model', 2, true);
			
			expect(isGradioTool(publicTool1)).toBe(true);
			expect(isGradioTool(publicTool2)).toBe(true);
			expect(isGradioTool(privateTool)).toBe(true);
		});

		it('should handle complex tool names with multiple underscores', () => {
			expect(isGradioTool('gr1_some_complex_tool_name_here')).toBe(true);
			expect(isGradioTool('grp5_another_complex_private_tool')).toBe(true);
		});
	});

	describe('should return false for non-Gradio tool names', () => {
		it('should reject standard HF tools', () => {
			expect(isGradioTool('hf_doc_search')).toBe(false);
			expect(isGradioTool('hf_model_search')).toBe(false);
			expect(isGradioTool('hf_whoami')).toBe(false);
			expect(isGradioTool('dynamic_space')).toBe(false);
		});

		it('should reject tools with missing number', () => {
			expect(isGradioTool('gr_tool')).toBe(false);
			expect(isGradioTool('grp_tool')).toBe(false);
		});

		it('should reject tools with missing underscore', () => {
			expect(isGradioTool('gr1tool')).toBe(false);
			expect(isGradioTool('grp2tool')).toBe(false);
		});

		it('should reject tools that do not start with gr/grp', () => {
			expect(isGradioTool('1_gr_tool')).toBe(false);
			expect(isGradioTool('some_gr1_tool')).toBe(false);
			expect(isGradioTool('prefix_grp2_tool')).toBe(false);
		});

		it('should reject empty or invalid inputs', () => {
			expect(isGradioTool('')).toBe(false);
			expect(isGradioTool('gr')).toBe(false);
			expect(isGradioTool('grp')).toBe(false);
			expect(isGradioTool('gr1')).toBe(false);
			expect(isGradioTool('grp1')).toBe(false);
		});

		it('should reject regular tool names', () => {
			expect(isGradioTool('regular_tool')).toBe(false);
			expect(isGradioTool('some_function')).toBe(false);
			expect(isGradioTool('tool_name')).toBe(false);
			expect(isGradioTool('api_call')).toBe(false);
		});

		it('should reject tools with invalid format variations', () => {
			expect(isGradioTool('gra1_tool')).toBe(false);
			expect(isGradioTool('grpp1_tool')).toBe(false);
			expect(isGradioTool('gr_1_tool')).toBe(false);
			expect(isGradioTool('grp_1_tool')).toBe(false);
		});
	});

	describe('edge cases', () => {
		it('should handle tools with numbers in the name part', () => {
			expect(isGradioTool('gr1_tool2')).toBe(true);
			expect(isGradioTool('grp1_v2_api')).toBe(true);
		});

		it('should handle tools with special characters in the name part', () => {
			expect(isGradioTool('gr1_tool-name')).toBe(true);
			expect(isGradioTool('grp1_tool.api')).toBe(true);
		});

		it('should require at least one digit', () => {
			expect(isGradioTool('gr_tool')).toBe(false);
			expect(isGradioTool('grp_tool')).toBe(false);
		});
	});
});

describe('createGradioToolName', () => {
	describe('should generate correct tool names', () => {
		it('should create public tool names with gr prefix', () => {
			expect(createGradioToolName('flux1_schnell', 0, false)).toBe('gr1_flux1_schnell');
			expect(createGradioToolName('model', 5, false)).toBe('gr6_model');
		});

		it('should create private tool names with grp prefix', () => {
			expect(createGradioToolName('private-model', 0, true)).toBe('grp1_private_model');
			expect(createGradioToolName('secret', 2, true)).toBe('grp3_secret');
		});

		it('should sanitize tool names correctly', () => {
			// Initial sanitization still happens (spaces, dots, dashes -> underscores)
			// But no underscore normalization
			expect(createGradioToolName('model-name.v2', 0, false)).toBe('gr1_model_name_v2');
			expect(createGradioToolName('my space test', 1, false)).toBe('gr2_my_space_test');
			expect(createGradioToolName('multi--dash__test', 0, false)).toBe('gr1_multi_dash__test');
		});

		it('should handle edge cases', () => {
			expect(createGradioToolName('', 0, false)).toBe('gr1_unknown');
			expect(createGradioToolName('simple', 0, false)).toBe('gr1_simple');
			expect(createGradioToolName('UPPERCASE-Model', 0, false)).toBe('gr1_uppercase_model');
		});

		it('should convert zero-based index to one-based tool names', () => {
			expect(createGradioToolName('test', 0, false)).toBe('gr1_test');
			expect(createGradioToolName('test', 1, false)).toBe('gr2_test');
			expect(createGradioToolName('test', 10, true)).toBe('grp11_test');
		});

		it('should enforce 49-character limit with middle truncation', () => {
			// Test basic truncation
			const longName = 'very_long_tool_name_that_exceeds_forty_nine_characters_total_and_more';
			const result1 = createGradioToolName(longName, 0, false);
			expect(result1.length).toBe(49);
			expect(result1).toBe('gr1_very_long_tool_name__haracters_total_and_more');
			
			// Test with larger index (less room for name)
			const result2 = createGradioToolName(longName, 999, true);
			expect(result2.length).toBe(49);
			expect(result2).toBe('grp1000_very_long_tool_name__cters_total_and_more');
			
			// Test that ending is preserved
			const toolWithUniqueEnd = 'common_prefix_for_tool_with_very_unique_ending_xyz123';
			const result3 = createGradioToolName(toolWithUniqueEnd, 0, false);
			expect(result3.length).toBe(49);
			expect(result3).toBe('gr1_common_prefix_for_to_ery_unique_ending_xyz123');
			expect(result3.endsWith('_xyz123')).toBe(true);
		});

		it('should handle edge cases with exactly 49 characters', () => {
			// Create a name that will result in exactly 49 chars: gr1_ + 45 chars = 49 total
			const exactName = 'tool_name_that_is_exactly_forty_five_characte'; // 45 chars
			const result = createGradioToolName(exactName, 0, false);
			expect(result).toBe('gr1_tool_name_that_is_exactly_forty_five_characte');
			expect(result.length).toBe(49);
		});

		it('should preserve unique endings when truncating', () => {
			// Names with same prefix but different endings - need longer names to trigger truncation at 49
			const tool1 = 'image_generation_model_for_advanced_processing_version_1_final';
			const tool2 = 'image_generation_model_for_advanced_processing_version_2_final';
			const result1 = createGradioToolName(tool1, 0, false);
			const result2 = createGradioToolName(tool2, 0, false);
			
			// Both should be truncated but keep different endings
			expect(result1).toBe('gr1_image_generation_mod_ocessing_version_1_final');
			expect(result2).toBe('gr1_image_generation_mod_ocessing_version_2_final');
			expect(result1).not.toBe(result2);
		});

		it('should not collide when one name is at limit and another exceeds by one char', () => {
			// Test case where one name is exactly at the limit and another is one char over
			// With index 29 (becomes "gr30"), we have 44 chars available for the name
			const tool1 = 'image_utilities_mcp_update_text_image_______'; // 44 chars - exactly at limit
			const tool2 = 'image_utilities_mcp_update_text_image________'; // 45 chars - exceeds by 1
			const result1 = createGradioToolName(tool1, 29, false);
			const result2 = createGradioToolName(tool2, 29, false);
			
			// No normalization for tool1, tool2 is truncated
			expect(result1).toBe('gr30_image_utilities_mcp_update_text_image_______');
			expect(result2).toBe('gr30_image_utilities_mcp__date_text_image________');
			expect(result1).not.toBe(result2);
			
			// First keeps its length, second is truncated to exactly 49
			expect(result1.length).toBe(49);
			expect(result2.length).toBe(49);
		});

		it('should prepend tool index to truncated names when provided', () => {
			// Test that toolIndex is prepended when truncation occurs
			const longName = 'very_long_tool_name_that_exceeds_forty_nine_characters_total_and_more';
			
			// Without toolIndex
			const result1 = createGradioToolName(longName, 0, false);
			expect(result1).toBe('gr1_very_long_tool_name__haracters_total_and_more');
			expect(result1.length).toBe(49);
			
			// With toolIndex 0
			const result2 = createGradioToolName(longName, 0, false, 0);
			expect(result2).toBe('gr1_0_very_long_tool_name__racters_total_and_more');
			expect(result2.length).toBe(49);
			
			// With toolIndex 1
			const result3 = createGradioToolName(longName, 0, false, 1);
			expect(result3).toBe('gr1_1_very_long_tool_name__racters_total_and_more');
			expect(result3.length).toBe(49);
			
			// Different tools should have different names
			expect(result2).not.toBe(result3);
		});

		it('should handle similar tool names with different indices correctly', () => {
			// Test multiple similar tool names that would collide without toolIndex
			const baseName = 'image_utilities_mcp_update_text_image________';
			
			const tools = [0, 1, 2, 3].map(toolIdx => 
				createGradioToolName(baseName, 29, false, toolIdx)
			);
			
			// All should be unique
			const uniqueTools = new Set(tools);
			expect(uniqueTools.size).toBe(4);
			
			// Each should start with its tool index after the prefix
			expect(tools[0]).toBe('gr30_0_image_utilities_mcp__te_text_image________');
			expect(tools[1]).toBe('gr30_1_image_utilities_mcp__te_text_image________');
			expect(tools[2]).toBe('gr30_2_image_utilities_mcp__te_text_image________');
			expect(tools[3]).toBe('gr30_3_image_utilities_mcp__te_text_image________');
		});

		it('should not add toolIndex when not truncating', () => {
			// Short names should not get toolIndex appended
			const shortName = 'simple_tool';
			
			const result1 = createGradioToolName(shortName, 0, false, 0);
			const result2 = createGradioToolName(shortName, 0, false, 1);
			
			// Both should be the same since no truncation happened
			expect(result1).toBe('gr1_simple_tool');
			expect(result2).toBe('gr1_simple_tool');
		});
	});

	describe('integration with tool names from endpoints', () => {
		it('should generate correct names for known tool names', () => {
			// These should match the examples of tool names from Gradio endpoints
			expect(createGradioToolName('flux1_schnell', 0, false)).toBe('gr1_flux1_schnell');
			expect(createGradioToolName('EasyGhibli', 1, false)).toBe('gr2_easyghibli');
		});
		});
});

describe('parseGradioSpaceIds', () => {
	describe('valid space IDs', () => {
		it('should parse single space ID', () => {
			const result = parseGradioSpaceIds('microsoft/Florence-2-large');
			expect(result).toEqual([{ name: 'microsoft/Florence-2-large' }]);
		});

		it('should parse multiple space IDs', () => {
			const result = parseGradioSpaceIds('microsoft/Florence-2-large,acme/foo,bar/baz');
			expect(result).toEqual([
				{ name: 'microsoft/Florence-2-large' },
				{ name: 'acme/foo' },
				{ name: 'bar/baz' },
			]);
		});

		it('should handle spaces around commas', () => {
			const result = parseGradioSpaceIds('microsoft/Florence-2-large, acme/foo , bar/baz');
			expect(result).toEqual([
				{ name: 'microsoft/Florence-2-large' },
				{ name: 'acme/foo' },
				{ name: 'bar/baz' },
			]);
		});

		it('should preserve case in space IDs', () => {
			const result = parseGradioSpaceIds('Microsoft/Florence-2-LARGE');
			expect(result).toEqual([{ name: 'Microsoft/Florence-2-LARGE' }]);
		});

		it('should handle spaces with special characters', () => {
			const result = parseGradioSpaceIds('user/space-name.v2,org/model_test');
			expect(result).toEqual([
				{ name: 'user/space-name.v2' },
				{ name: 'org/model_test' },
			]);
		});

		it('should handle real-world example from bug report', () => {
			// This was the actual failing case: microsoft/Florence-2-large
			const result = parseGradioSpaceIds('microsoft/Florence-2-large');
			expect(result).toEqual([{ name: 'microsoft/Florence-2-large' }]);
			// Verify we're NOT converting the slash to a dash here
			expect(result[0].name).not.toContain('-Florence-'); // Should have /Florence-
			expect(result[0].name).toContain('/Florence-');
		});
	});

	describe('special sentinel values', () => {
		it('should return empty array for "none"', () => {
			const result = parseGradioSpaceIds('none');
			expect(result).toEqual([]);
		});

		it('should return empty array for "NONE" (case insensitive)', () => {
			const result = parseGradioSpaceIds('NONE');
			expect(result).toEqual([]);
		});

		it('should return empty array for "None"', () => {
			const result = parseGradioSpaceIds('None');
			expect(result).toEqual([]);
		});

		it('should skip "none" in comma-separated list', () => {
			const result = parseGradioSpaceIds('microsoft/Florence-2-large,none,acme/foo');
			expect(result).toEqual([
				{ name: 'microsoft/Florence-2-large' },
				{ name: 'acme/foo' },
			]);
		});

		it('should handle multiple "none" values', () => {
			const result = parseGradioSpaceIds('none,none,acme/foo,none');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});
	});

	describe('empty and whitespace inputs', () => {
		it('should handle empty string', () => {
			const result = parseGradioSpaceIds('');
			expect(result).toEqual([]);
		});

		it('should handle whitespace-only string', () => {
			const result = parseGradioSpaceIds('   ');
			expect(result).toEqual([]);
		});

		it('should handle string with only commas', () => {
			const result = parseGradioSpaceIds(',,,');
			expect(result).toEqual([]);
		});

		it('should filter out empty entries between commas', () => {
			const result = parseGradioSpaceIds('acme/foo,,bar/baz');
			expect(result).toEqual([
				{ name: 'acme/foo' },
				{ name: 'bar/baz' },
			]);
		});
	});

	describe('invalid space IDs', () => {
		it('should skip entries with no slash', () => {
			const result = parseGradioSpaceIds('invalid-space,acme/foo');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});

		it('should skip entries with multiple slashes', () => {
			const result = parseGradioSpaceIds('microsoft/spaces/Florence-2-large,acme/foo');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});

		it('should skip entries with trailing slash', () => {
			const result = parseGradioSpaceIds('microsoft/,acme/foo');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});

		it('should skip entries with leading slash', () => {
			const result = parseGradioSpaceIds('/Florence-2-large,acme/foo');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});

		it('should skip entries with only a slash', () => {
			const result = parseGradioSpaceIds('/,acme/foo');
			expect(result).toEqual([{ name: 'acme/foo' }]);
		});

		it('should handle mix of valid and invalid entries', () => {
			const result = parseGradioSpaceIds('valid/space,invalid,another/valid,too/many/slashes');
			expect(result).toEqual([
				{ name: 'valid/space' },
				{ name: 'another/valid' },
			]);
		});
	});

	describe('URL encoding scenarios', () => {
		it('should handle already-decoded slash (Express behavior)', () => {
			// In real usage, Express decodes %2F to / before our code sees it
			const result = parseGradioSpaceIds('microsoft/Florence-2-large');
			expect(result).toEqual([{ name: 'microsoft/Florence-2-large' }]);
		});

		it('should handle multiple spaces with decoded slashes', () => {
			// Client sends: microsoft%2FFoo,acme%2Fbar
			// Express decodes to: microsoft/Foo,acme/bar
			const result = parseGradioSpaceIds('microsoft/Foo,acme/bar');
			expect(result).toEqual([
				{ name: 'microsoft/Foo' },
				{ name: 'acme/bar' },
			]);
		});
	});

	describe('edge cases from bug investigation', () => {
		it('should NOT manually construct subdomains', () => {
			// The old buggy code did: entry.replace(/[/.]/g, '-')
			// This test ensures we DON'T do that transformation
			const result = parseGradioSpaceIds('microsoft/Florence-2-large');

			// We should preserve the original space ID exactly
			expect(result[0].name).toBe('microsoft/Florence-2-large');

			// Verify we're not returning a subdomain at this stage
			expect(result[0]).not.toHaveProperty('subdomain');
		});

		it('should preserve dots in space names', () => {
			const result = parseGradioSpaceIds('user/model.v2');
			expect(result[0].name).toBe('user/model.v2');
			expect(result[0].name).toContain('.');
		});

		it('should preserve underscores in space names', () => {
			const result = parseGradioSpaceIds('user/my_model');
			expect(result[0].name).toBe('user/my_model');
			expect(result[0].name).toContain('_');
		});
	});
});

describe('resolveMcpFileListingSource', () => {
	const DYNAMIC_SPACE_TOOL_ID = 'dynamic_space';
	const baseParams = {
		username: 'alice',
		token: 'hf_token',
		gradioSpaceCount: 1,
		builtInTools: [],
		dynamicSpaceToolId: DYNAMIC_SPACE_TOOL_ID,
	};

	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('returns legacy gradio-files dataset when it exists', async () => {
		vi.mocked(hub.repoExists).mockResolvedValueOnce(true);

		await expect(resolveMcpFileListingSource(baseParams)).resolves.toEqual({
			kind: 'dataset',
			id: 'alice/gradio-files',
		});
		expect(hub.repoExists).toHaveBeenCalledTimes(1);
		expect(hub.repoExists).toHaveBeenCalledWith({
			repo: { type: 'dataset', name: 'alice/gradio-files' },
			accessToken: 'hf_token',
		});
	});

	it('returns null when the legacy dataset is missing', async () => {
		vi.mocked(hub.repoExists).mockResolvedValueOnce(false);

		await expect(resolveMcpFileListingSource(baseParams)).resolves.toBeNull();
	});

	it('returns null without auth prerequisites', async () => {
		await expect(resolveMcpFileListingSource({ ...baseParams, token: undefined })).resolves.toBeNull();
		await expect(resolveMcpFileListingSource({ ...baseParams, username: undefined })).resolves.toBeNull();
		expect(hub.repoExists).not.toHaveBeenCalled();
	});

	it('returns null when file listing is not relevant', async () => {
		await expect(
			resolveMcpFileListingSource({
				...baseParams,
				gradioSpaceCount: 0,
				builtInTools: [],
			})
		).resolves.toBeNull();
		expect(hub.repoExists).not.toHaveBeenCalled();
	});

	it('allows dynamic_space to make file listing relevant', async () => {
		vi.mocked(hub.repoExists).mockResolvedValueOnce(true);

		await expect(
			resolveMcpFileListingSource({
				...baseParams,
				gradioSpaceCount: 0,
				builtInTools: [DYNAMIC_SPACE_TOOL_ID],
			})
		).resolves.toEqual({ kind: 'dataset', id: 'alice/gradio-files' });
	});
});