Skip to content

Commit b1d8674

Browse files
epessinaepessina
committed
add AI feature check in ITD list
1 parent fc19f55 commit b1d8674

File tree

3 files changed

+251
-5
lines changed

3 files changed

+251
-5
lines changed

‎src/tools/descriptions.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ Do never attemp to fetch all the logs without a date range filter, as it could l
7171
`,
7272
LIST_MARKETPLACE_ITEMS_VERSIONS: 'List all the available versions of a marketplace item',
7373
MARKETPLACE_ITEM_VERSION_INFO: 'Get information about a specific version of a marketplace item',
74-
LIST_MARKETPLACE_ITEM_TYPE_DEFINITIONS: 'List all the marketplace Item Type Definitions the caller has permission to see (i.e., the ones available to all tenants and the private ones of tenants the user has permission to see)',
74+
LIST_MARKETPLACE_ITEM_TYPE_DEFINITIONS: 'List the metadata of all the marketplace Item Type Definitions the caller has permission to see (i.e., the ones available to all tenants and the private ones of tenants the user has permission to see)',
7575
MARKETPLACE_ITEM_TYPE_DEFINITION_INFO: 'Get information about a specific Item Type Definition identified by its compound primary key as path parameters (i.e., id of the tenant namespace, and name of the definition)',
7676

7777
// project tools

‎src/tools/marketplace/index.test.ts‎

Lines changed: 195 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,19 @@ const itemTypeDefinitions: CatalogItemTypeDefinition[] = [
7373
displayName: 'Plugin ITD',
7474
name: 'plugin',
7575
visibility: { scope: 'console' },
76+
annotations: { foo: 'bar' },
77+
description: 'Description...',
78+
documentation: { url: 'http://example.com', type: 'external' },
79+
icon: { base64Data: 'abc', mediaType: 'image/png' },
80+
labels: { foo: 'bar' },
81+
links: [ { url: 'example.com', displayName: 'Example' } ],
82+
maintainers: [ { email: '[email protected]', name: 'John Doe' } ],
83+
publisher: {
84+
name: 'John Doe',
85+
url: 'http://example.com',
86+
image: { base64Data: 'abc', mediaType: 'image/png' },
87+
},
88+
tags: [ 'foo', 'bar' ],
7689
},
7790
spec: {
7891
isVersioningSupported: true,
@@ -413,7 +426,163 @@ suite('marketplace Item Type Definitions list tool', async () => {
413426
})
414427

415428
test('should return ITDs', async (t) => {
429+
const isAiFeaturesEnabledForTenantMockFn = mock.fn(async () => true)
430+
431+
const client = await getTestMCPServerClient({
432+
isAiFeaturesEnabledForTenantMockFn,
433+
listMarketplaceItemTypeDefinitionsMockFn,
434+
})
435+
436+
const result = await client.request({
437+
method: 'tools/call',
438+
params: {
439+
name: 'list_marketplace_item_type_definitions',
440+
arguments: {},
441+
},
442+
}, CallToolResultSchema)
443+
444+
const expectedITDs = [
445+
{
446+
annotations: { foo: 'bar' },
447+
description: 'Description...',
448+
displayName: 'Plugin ITD',
449+
documentation: { url: 'http://example.com', type: 'external' },
450+
labels: { foo: 'bar' },
451+
links: [ { url: 'example.com', displayName: 'Example' } ],
452+
maintainers: [ { email: '[email protected]', name: 'John Doe' } ],
453+
name: 'plugin',
454+
namespace: { scope: 'tenant', id: 'mia-platform' },
455+
publisher: {
456+
name: 'John Doe',
457+
url: 'http://example.com',
458+
},
459+
tags: [ 'foo', 'bar' ],
460+
visibility: { scope: 'console' },
461+
},
462+
{
463+
displayName: 'Custom Workload ITD',
464+
name: 'custom-workload',
465+
namespace: { scope: 'tenant', id: 'my-company' },
466+
visibility: { scope: 'tenant', ids: [ 'my-company' ] },
467+
},
468+
]
469+
470+
t.assert.deepEqual(result.content, [
471+
{
472+
text: JSON.stringify(expectedITDs),
473+
type: 'text',
474+
},
475+
])
476+
477+
t.assert.equal(isAiFeaturesEnabledForTenantMockFn.mock.callCount(), 2)
478+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(0)!.arguments, [ 'mia-platform' ])
479+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(1)!.arguments, [ 'my-company' ])
480+
})
481+
482+
test('should return ITDs with namespace param if all namespaces have AI features enabled', async (t) => {
483+
const isAiFeaturesEnabledForTenantMockFn = mock.fn(async () => true)
484+
485+
const client = await getTestMCPServerClient({
486+
isAiFeaturesEnabledForTenantMockFn,
487+
listMarketplaceItemTypeDefinitionsMockFn,
488+
})
489+
490+
const result = await client.request({
491+
method: 'tools/call',
492+
params: {
493+
name: 'list_marketplace_item_type_definitions',
494+
arguments: {
495+
namespace: 'mia-platform,my-company,other-company',
496+
},
497+
},
498+
}, CallToolResultSchema)
499+
500+
const expectedITDs = [
501+
{
502+
annotations: { foo: 'bar' },
503+
description: 'Description...',
504+
displayName: 'Plugin ITD',
505+
documentation: { url: 'http://example.com', type: 'external' },
506+
labels: { foo: 'bar' },
507+
links: [ { url: 'example.com', displayName: 'Example' } ],
508+
maintainers: [ { email: '[email protected]', name: 'John Doe' } ],
509+
name: 'plugin',
510+
namespace: { scope: 'tenant', id: 'mia-platform' },
511+
publisher: {
512+
name: 'John Doe',
513+
url: 'http://example.com',
514+
},
515+
tags: [ 'foo', 'bar' ],
516+
visibility: { scope: 'console' },
517+
},
518+
{
519+
displayName: 'Custom Workload ITD',
520+
name: 'custom-workload',
521+
namespace: { scope: 'tenant', id: 'my-company' },
522+
visibility: { scope: 'tenant', ids: [ 'my-company' ] },
523+
},
524+
]
525+
526+
t.assert.deepEqual(result.content, [
527+
{
528+
text: JSON.stringify(expectedITDs),
529+
type: 'text',
530+
},
531+
])
532+
533+
t.assert.equal(isAiFeaturesEnabledForTenantMockFn.mock.callCount(), 3)
534+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(0)!.arguments, [ 'mia-platform' ])
535+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(1)!.arguments, [ 'my-company' ])
536+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(2)!.arguments, [ 'other-company' ])
537+
})
538+
539+
test('should return error message if namespace param is used and not all namespaces have AI features enabled', async (t) => {
540+
const isAiFeaturesEnabledForTenantMockFn = mock.fn(async (tenantId: string) => {
541+
if (tenantId === 'other-company') {
542+
throw new Error('error message')
543+
}
544+
545+
return true
546+
})
547+
416548
const client = await getTestMCPServerClient({
549+
isAiFeaturesEnabledForTenantMockFn,
550+
listMarketplaceItemTypeDefinitionsMockFn,
551+
})
552+
553+
const result = await client.request({
554+
method: 'tools/call',
555+
params: {
556+
name: 'list_marketplace_item_type_definitions',
557+
arguments: {
558+
namespace: 'mia-platform,other-company',
559+
},
560+
},
561+
}, CallToolResultSchema)
562+
563+
t.assert.deepEqual(result.content, [
564+
{
565+
text: 'Error fetching marketplace item type definitions: error message',
566+
type: 'text',
567+
},
568+
])
569+
570+
t.assert.equal(isAiFeaturesEnabledForTenantMockFn.mock.callCount(), 2)
571+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(0)!.arguments, [ 'mia-platform' ])
572+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(1)!.arguments, [ 'other-company' ])
573+
})
574+
575+
test('should return ITDs filtering out namespaces without AI features enabled', async (t) => {
576+
const isAiFeaturesEnabledForTenantMockFn = mock.fn(async (tenantId: string) => {
577+
if (tenantId === 'my-company') {
578+
throw new Error('error message')
579+
}
580+
581+
return true
582+
})
583+
584+
const client = await getTestMCPServerClient({
585+
isAiFeaturesEnabledForTenantMockFn,
417586
listMarketplaceItemTypeDefinitionsMockFn,
418587
})
419588

@@ -425,16 +594,41 @@ suite('marketplace Item Type Definitions list tool', async () => {
425594
},
426595
}, CallToolResultSchema)
427596

597+
const expectedITDs = [
598+
{
599+
annotations: { foo: 'bar' },
600+
description: 'Description...',
601+
displayName: 'Plugin ITD',
602+
documentation: { url: 'http://example.com', type: 'external' },
603+
labels: { foo: 'bar' },
604+
links: [ { url: 'example.com', displayName: 'Example' } ],
605+
maintainers: [ { email: '[email protected]', name: 'John Doe' } ],
606+
name: 'plugin',
607+
namespace: { scope: 'tenant', id: 'mia-platform' },
608+
publisher: {
609+
name: 'John Doe',
610+
url: 'http://example.com',
611+
},
612+
tags: [ 'foo', 'bar' ],
613+
visibility: { scope: 'console' },
614+
},
615+
]
616+
428617
t.assert.deepEqual(result.content, [
429618
{
430-
text: JSON.stringify(itemTypeDefinitions),
619+
text: JSON.stringify(expectedITDs),
431620
type: 'text',
432621
},
433622
])
623+
624+
t.assert.equal(isAiFeaturesEnabledForTenantMockFn.mock.callCount(), 2)
625+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(0)!.arguments, [ 'mia-platform' ])
626+
t.assert.deepEqual(isAiFeaturesEnabledForTenantMockFn.mock.calls.at(1)!.arguments, [ 'my-company' ])
434627
})
435628

436629
test('should return error message if request return error', async (t) => {
437630
const client = await getTestMCPServerClient({
631+
isAiFeaturesEnabledForTenantMockFn: async () => true,
438632
listMarketplaceItemTypeDefinitionsMockFn,
439633
})
440634

‎src/tools/marketplace/index.ts‎

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
1818
import { z } from 'zod'
1919

2020
import { assertAiFeaturesEnabledForTenant } from '../utils/validations'
21+
import { CatalogItemTypeDefinition } from '@mia-platform/console-types'
2122
import { CatalogItemTypes } from '../../apis/types/marketplace'
2223
import { IAPIClient } from '../../apis/client'
2324
import { paramsDescriptions, toolNames, toolsDescriptions } from '../descriptions'
@@ -158,16 +159,67 @@ export function addMarketplaceCapabilities (server: McpServer, client: IAPIClien
158159
},
159160
async ({ namespace, name, displayName }): Promise<CallToolResult> => {
160161
try {
161-
// TODO: check namespace AI validity
162+
if (namespace) {
163+
const tenantIds = namespace.split(',')
164+
await Promise.all(tenantIds.map((tenantId) => assertAiFeaturesEnabledForTenant(client, tenantId)))
165+
}
162166

163-
// TODO: should I map the data to take only a subset of fields
164167
const data = await client.listMarketplaceItemTypeDefinitions(namespace, name, displayName)
165168

169+
let tenantIdAiFeaturesEnabledMap: Map<string, boolean> | null = null
170+
171+
// We don't need to check the returned data if the argument is used, since we already checked that all
172+
// passed namespaces have AI features enabled
173+
if (!namespace) {
174+
tenantIdAiFeaturesEnabledMap = new Map()
175+
176+
const dataTenantIds = new Set(data.map((itd) => itd.metadata.namespace.id))
177+
178+
const assertAiFeaturesEnabledForTenantPromises: Promise<void>[] = []
179+
dataTenantIds.forEach((tenantId) => {
180+
assertAiFeaturesEnabledForTenantPromises.push(assertAiFeaturesEnabledForTenant(client, tenantId).
181+
then(() => {
182+
tenantIdAiFeaturesEnabledMap?.set(tenantId, true)
183+
}).
184+
catch(() => {
185+
tenantIdAiFeaturesEnabledMap?.set(tenantId, false)
186+
}))
187+
})
188+
189+
await Promise.all(assertAiFeaturesEnabledForTenantPromises)
190+
}
191+
192+
const mappedData = data.reduce<CatalogItemTypeDefinition['metadata'][]>((acc, itd) => {
193+
if (tenantIdAiFeaturesEnabledMap?.get(itd.metadata.namespace.id) === false) {
194+
return acc
195+
}
196+
197+
return [
198+
...acc,
199+
{
200+
annotations: itd.metadata.annotations,
201+
description: itd.metadata.description,
202+
displayName: itd.metadata.displayName,
203+
documentation: itd.metadata.documentation,
204+
labels: itd.metadata.labels,
205+
links: itd.metadata.links,
206+
maintainers: itd.metadata.maintainers,
207+
name: itd.metadata.name,
208+
namespace: itd.metadata.namespace,
209+
publisher: itd.metadata.publisher
210+
? { name: itd.metadata.publisher.name, url: itd.metadata.publisher.url }
211+
: undefined,
212+
tags: itd.metadata.tags,
213+
visibility: itd.metadata.visibility,
214+
},
215+
]
216+
}, [])
217+
166218
return {
167219
content: [
168220
{
169221
type: 'text',
170-
text: JSON.stringify(data),
222+
text: JSON.stringify(mappedData),
171223
},
172224
],
173225
}

0 commit comments

Comments
 (0)