diff --git a/apps/cloudflare-one-casb/src/tools/integrations.ts b/apps/cloudflare-one-casb/src/tools/integrations.ts index 0d8072a0..38a6c68b 100644 --- a/apps/cloudflare-one-casb/src/tools/integrations.ts +++ b/apps/cloudflare-one-casb/src/tools/integrations.ts @@ -16,34 +16,51 @@ import { assetCategoryVendorParam, } from '@repo/mcp-common/src/schemas/cf1-integrations' +import type { zReturnedAssetResult } from '@repo/mcp-common/src/schemas/cf1-integrations' import type { ToolDefinition } from '@repo/mcp-common/src/types/tools' import type { CASBMCP } from '../index' -const PAGE_SIZE = 3 - const integrationIdParam = z.string().describe('The UUID of the integration to analyze') const assetSearchTerm = z.string().describe('The search keyword for assets') const assetIdParam = z.string().describe('The UUID of the asset to analyze') const assetCategoryIdParam = z.string().describe('The UUID of the asset category to analyze') +const paginationParams = z + .object({ + page: z.number().optional().describe('Page number for pagination'), + pageSize: z.number().optional().describe('Number of items per page'), + }) + .optional() + +const assetLite = ({ fields: _, integration, ...asset }: zReturnedAssetResult) => ({ + ...asset, + integration: { + id: integration.id, + name: integration.name, + }, +}) + const toolDefinitions: Array> = [ { name: 'integration_by_id', description: 'Analyze Cloudflare One Integration by ID', - params: { integrationIdParam }, + params: { integrationIdParam, paginationParams }, handler: async ({ integrationIdParam, accountId, apiToken, + paginationParams = {}, }: { integrationIdParam: string accountId: string apiToken: string + paginationParams?: { page?: number; pageSize?: number } }) => { const { integration } = await handleIntegrationById({ integrationIdParam, accountId, apiToken, + ...paginationParams, }) return { integration } }, @@ -51,51 +68,68 @@ const toolDefinitions: Array> = [ { name: 'integrations_list', description: 'List all Cloudflare One Integrations in a given account', - params: {}, - handler: async ({ accountId, apiToken }: { accountId: string; apiToken: string }) => { - const { integrations } = await handleIntegrations({ accountId, apiToken }) - return { integrations } + params: { paginationParams }, + handler: async ({ + accountId, + apiToken, + paginationParams = {}, + }: { + accountId: string + apiToken: string + paginationParams?: { page?: number; pageSize?: number } + }) => { + const { integrations, result_info } = await handleIntegrations({ + accountId, + apiToken, + ...paginationParams, + }) + return { integrations, result_info } }, }, { name: 'assets_search', description: 'Search Assets by keyword', - params: { assetSearchTerm }, + params: { assetSearchTerm, paginationParams }, handler: async ({ assetSearchTerm, accountId, apiToken, + paginationParams = {}, }: { assetSearchTerm: string accountId: string apiToken: string + paginationParams?: { page?: number; pageSize?: number } }) => { - const { assets } = await handleAssetsSearch({ + const { assets, result_info } = await handleAssetsSearch({ accountId, apiToken, searchTerm: assetSearchTerm, - pageSize: PAGE_SIZE, + ...paginationParams, }) - return { assets } + return { assets: assets.map(assetLite), result_info } }, }, { name: 'asset_by_id', description: 'Search Assets by ID', - params: { assetIdParam }, + params: { assetIdParam, paginationParams }, handler: async ({ assetIdParam, accountId, apiToken, + paginationParams = {}, }: { assetIdParam: string accountId: string apiToken: string + paginationParams?: { page?: number; pageSize?: number } }) => { const { asset } = await handleAssetById({ accountId, apiToken, assetId: assetIdParam, + ...paginationParams, }) return { asset } }, @@ -103,144 +137,162 @@ const toolDefinitions: Array> = [ { name: 'assets_by_integration_id', description: 'Search Assets by Integration ID', - params: { integrationIdParam }, + params: { integrationIdParam, paginationParams }, handler: async ({ integrationIdParam, accountId, apiToken, + paginationParams = {}, }: { integrationIdParam: string accountId: string apiToken: string + paginationParams?: { page?: number; pageSize?: number } }) => { - const { assets } = await handleAssetsByIntegrationId({ + const { assets, result_info } = await handleAssetsByIntegrationId({ accountId, apiToken, integrationId: integrationIdParam, - pageSize: PAGE_SIZE, + ...paginationParams, }) - return { assets } + return { assets: assets.map(assetLite), result_info } }, }, { name: 'assets_by_category_id', description: 'Search Assets by Asset Category ID', - params: { assetCategoryIdParam }, + params: { assetCategoryIdParam, paginationParams }, handler: async ({ assetCategoryIdParam, accountId, apiToken, + paginationParams = {}, }: { assetCategoryIdParam: string accountId: string apiToken: string + paginationParams?: { page?: number; pageSize?: number } }) => { - const { assets } = await handleAssetsByAssetCategoryId({ + const { assets, result_info } = await handleAssetsByAssetCategoryId({ accountId, apiToken, categoryId: assetCategoryIdParam, - pageSize: PAGE_SIZE, + ...paginationParams, }) - return { assets } + return { assets: assets.map(assetLite), result_info } }, }, { name: 'assets_list', description: 'Paginated list of Assets', - params: {}, - handler: async ({ accountId, apiToken }: { accountId: string; apiToken: string }) => { - const { assets } = await handleAssets({ + params: { paginationParams }, + handler: async ({ + accountId, + apiToken, + paginationParams = {}, + }: { + accountId: string + apiToken: string + paginationParams?: { page?: number; pageSize?: number } + }) => { + const { assets, result_info } = await handleAssets({ accountId, apiToken, - pageSize: PAGE_SIZE, + ...paginationParams, }) - return { assets } + return { assets: assets.map(assetLite), result_info } }, }, { name: 'asset_categories_list', description: 'List Asset Categories', - params: {}, - handler: async ({ accountId, apiToken }: { accountId: string; apiToken: string }) => { - const { categories } = await handleAssetCategories({ + params: { paginationParams }, + handler: async ({ + accountId, + apiToken, + paginationParams = {}, + }: { + accountId: string + apiToken: string + paginationParams?: { page?: number; pageSize?: number } + }) => + await handleAssetCategories({ accountId, apiToken, - }) - return { categories } - }, + ...paginationParams, + }), }, { name: 'asset_categories_by_vendor', description: 'List asset categories by vendor', - params: { assetCategoryVendorParam }, + params: { assetCategoryVendorParam, paginationParams }, handler: async ({ assetCategoryVendorParam, accountId, apiToken, + paginationParams = {}, }: { assetCategoryVendorParam: string accountId: string apiToken: string - }) => { - const { categories } = await handleAssetCategories({ + paginationParams?: { page?: number; pageSize?: number } + }) => + await handleAssetCategories({ accountId, apiToken, vendor: assetCategoryVendorParam, - }) - return { categories } - }, + ...paginationParams, + }), }, { name: 'asset_categories_by_type', description: 'Search Asset Categories by type', - params: { assetCategoryTypeParam }, + params: { assetCategoryTypeParam, paginationParams }, handler: async ({ assetCategoryTypeParam, accountId, apiToken, + paginationParams = {}, }: { assetCategoryTypeParam?: string accountId: string apiToken: string - }) => { - const { categories } = await handleAssetCategories({ + paginationParams?: { page?: number; pageSize?: number } + }) => + await handleAssetCategories({ accountId, apiToken, type: assetCategoryTypeParam, - }) - return { categories } - }, + ...paginationParams, + }), }, { name: 'asset_categories_by_vendor_and_type', description: 'Search Asset Categories by vendor and type', - params: { assetCategoryTypeParam, assetCategoryVendorParam }, + params: { assetCategoryTypeParam, assetCategoryVendorParam, paginationParams }, handler: async ({ assetCategoryTypeParam, assetCategoryVendorParam, accountId, apiToken, + paginationParams = {}, }: { assetCategoryTypeParam?: string assetCategoryVendorParam: string accountId: string apiToken: string - }) => { - const { categories } = await handleAssetCategories({ + paginationParams?: { page?: number; pageSize?: number } + }) => + await handleAssetCategories({ accountId, apiToken, type: assetCategoryTypeParam, vendor: assetCategoryVendorParam, - }) - return { categories } - }, + ...paginationParams, + }), }, ] -/** - * Registers the logs analysis tool with the MCP server - * @param agent The MCP server instance - */ export function registerIntegrationsTools(agent: CASBMCP) { toolDefinitions.forEach(({ name, description, params, handler }) => { agent.server.tool(name, description, params, withAccountCheck(agent, handler)) diff --git a/packages/mcp-common/src/api/cf1-integration.ts b/packages/mcp-common/src/api/cf1-integration.ts index 55a4b660..693f33d5 100644 --- a/packages/mcp-common/src/api/cf1-integration.ts +++ b/packages/mcp-common/src/api/cf1-integration.ts @@ -11,10 +11,14 @@ import { V4Schema } from '../v4-api' import type { z } from 'zod' import type { zReturnedAssetCategoriesResult, + zReturnedAssetResult, zReturnedAssetsResult, zReturnedIntegrationResult, zReturnedIntegrationsResult, } from '../schemas/cf1-integrations' +import type { zV4ResultInfoSchema } from '../v4-api' + +const DEFAULT_PAGE_SIZE = 3 interface BaseParams { accountId: string @@ -26,17 +30,20 @@ interface PaginationParams { pageSize?: number } -type IntegrationParams = BaseParams & { integrationIdParam: string } -type AssetCategoryParams = BaseParams & { type?: string; vendor?: string } +type IntegrationParams = BaseParams & { integrationIdParam: string } & PaginationParams +type AssetCategoryParams = BaseParams & { type?: string; vendor?: string } & PaginationParams type AssetSearchParams = BaseParams & { searchTerm: string } & PaginationParams -type AssetByIdParams = BaseParams & { assetId: string } +type AssetByIdParams = BaseParams & { assetId: string } & PaginationParams type AssetByCategoryParams = BaseParams & { categoryId: string } & PaginationParams type AssetByIntegrationParams = BaseParams & { integrationId: string } & PaginationParams -const buildParams = (baseParams: Record, pagination?: PaginationParams) => { +const buildParams = (baseParams: Record, pagination: PaginationParams = {}) => { const params = new URLSearchParams(baseParams) - if (pagination?.page) params.append('page', String(pagination.page)) - if (pagination?.pageSize) params.append('page_size', String(pagination.pageSize)) + const pageSize = pagination.pageSize ?? DEFAULT_PAGE_SIZE + + if (pagination.page) params.append('page', String(pagination.page)) + params.append('page_size', String(pageSize)) + return params } @@ -57,7 +64,7 @@ const makeApiCall = async ({ apiToken: string responseSchema: z.ZodType params?: URLSearchParams -}): Promise => { +}): Promise<{ result: T; result_info: zV4ResultInfoSchema }> => { try { const fullEndpoint = params ? `${endpoint}?${params.toString()}` : endpoint const data = await fetchCloudflareApi({ @@ -70,20 +77,20 @@ const makeApiCall = async ({ headers: { 'Content-Type': 'application/json' }, }, }) - return data.result as T + return data } catch (error) { console.error(`API call failed for ${endpoint}:`, error) throw error } } -// Resource-specific API call handlers const makeIntegrationCall = (params: IntegrationParams, responseSchema: z.ZodType) => makeApiCall({ endpoint: buildIntegrationEndpoint(params.integrationIdParam), accountId: params.accountId, apiToken: params.apiToken, responseSchema, + params: buildParams({}, params), }) const makeAssetCall = ( @@ -106,64 +113,74 @@ const makeAssetCategoryCall = (params: AssetCategoryParams, responseSchema: z accountId: params.accountId, apiToken: params.apiToken, responseSchema, - params: buildParams({ - ...(params.vendor && { vendor: params.vendor }), - ...(params.type && { type: params.type }), - }), + params: buildParams( + { + ...(params.vendor && { vendor: params.vendor }), + ...(params.type && { type: params.type }), + }, + params + ), }) -// Integration handlers export async function handleIntegrationById( params: IntegrationParams -): Promise<{ integration: zReturnedIntegrationResult | null }> { - const integration = await makeIntegrationCall( +): Promise<{ integration: zReturnedIntegrationResult | null; result_info: zV4ResultInfoSchema }> { + const data = await makeIntegrationCall( params, V4Schema(IntegrationResponse) ) - return { integration } + return { integration: data.result, result_info: data?.result_info } } export async function handleIntegrations( - params: BaseParams -): Promise<{ integrations: zReturnedIntegrationsResult | null }> { - const integrations = await makeApiCall({ + params: BaseParams & PaginationParams +): Promise<{ integrations: zReturnedIntegrationsResult | null; result_info: zV4ResultInfoSchema }> { + const data = await makeApiCall({ endpoint: '/casb/integrations', accountId: params.accountId, apiToken: params.apiToken, responseSchema: V4Schema(IntegrationsResponse), + params: buildParams({}, params), }) - return { integrations } + + return { integrations: data.result, result_info: data?.result_info } } -// Asset category handlers -export async function handleAssetCategories( - params: AssetCategoryParams -): Promise<{ categories: zReturnedAssetCategoriesResult | null }> { - const categories = await makeAssetCategoryCall( - params, - V4Schema(AssetCategoriesResponse) - ) - return { categories } +export async function handleAssetCategories(params: AssetCategoryParams): Promise<{ + categories: zReturnedAssetCategoriesResult | null + result_info: zV4ResultInfoSchema +}> { + const { result: categories, result_info } = + await makeAssetCategoryCall( + params, + V4Schema(AssetCategoriesResponse) + ) + + return { categories, result_info } } -// Asset handlers export async function handleAssets(params: BaseParams & PaginationParams) { - const assets = await makeAssetCall(params, V4Schema(AssetsResponse)) - return { assets } + const { result: assets, result_info } = await makeAssetCall( + params, + V4Schema(AssetsResponse) + ) + + return { assets, result_info } } export async function handleAssetsByIntegrationId(params: AssetByIntegrationParams) { - const assets = await makeAssetCall( + const { result: assets, result_info } = await makeAssetCall( params, V4Schema(AssetsResponse), undefined, { integration_id: params.integrationId } ) - return { assets } + + return { assets, result_info } } export async function handleAssetById(params: AssetByIdParams) { - const asset = await makeAssetCall( + const asset = await makeAssetCall( params, V4Schema(AssetDetail), params.assetId @@ -172,21 +189,24 @@ export async function handleAssetById(params: AssetByIdParams) { } export async function handleAssetsByAssetCategoryId(params: AssetByCategoryParams) { - const assets = await makeAssetCall( + const data = await makeAssetCall( params, V4Schema(AssetsResponse), undefined, { category_id: params.categoryId } ) - return { assets } + return { assets: data.result, result_info: data.result_info } } -export async function handleAssetsSearch(params: AssetSearchParams) { - const assets = await makeAssetCall( +export async function handleAssetsSearch( + params: AssetSearchParams +): Promise<{ assets: zReturnedAssetsResult; result_info: zV4ResultInfoSchema }> { + const data = await makeAssetCall( params, V4Schema(AssetsResponse), undefined, { search: params.searchTerm } ) - return { assets } + + return { assets: data.result, result_info: data?.result_info } } diff --git a/packages/mcp-common/src/v4-api.ts b/packages/mcp-common/src/v4-api.ts index 135c56b1..1025bdc5 100644 --- a/packages/mcp-common/src/v4-api.ts +++ b/packages/mcp-common/src/v4-api.ts @@ -8,6 +8,16 @@ const V4ErrorSchema = z.array( }) ) +const V4ResultInfoSchema = z.object({ + next: z.string().nullable(), + previous: z.string().nullable(), + per_page: z.number().nullable(), + count: z.number().nullable(), + total_count: z.number().nullable(), +}) + +export type zV4ResultInfoSchema = z.infer + export const V4Schema = ( resultType: TResultType ): z.ZodObject<{ @@ -17,6 +27,7 @@ export const V4Schema = ( messages: z.ZodArray }> => z.object({ + result_info: V4ResultInfoSchema.nullable(), result: resultType.nullable(), success: z.boolean(), errors: V4ErrorSchema,