integration/Plugins/BlueprintMCP/Source/BlueprintMCP/Private/MCPAssets.cpp

#include "MCPAssets.h"
#include "MCPServer.h"
#include "Engine/Blueprint.h"
#include "Engine/World.h"
#include "Engine/Level.h"
#include "Engine/LevelScriptBlueprint.h"
#include "Materials/Material.h"
#include "MCPUtils.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "AssetRegistry/IAssetRegistry.h"

// ============================================================
// MCPAssetsBase
// ============================================================

MCPAssetsBase::MCPAssetsBase(UClass* InTargetClass)
	: TargetClass(InTargetClass)
{
	Scans.Add(InTargetClass);
}

MCPAssetsBase& MCPAssetsBase::Exact(const FString& InName)
{
	MatchName = InName;
	bExactMatch = true;
	bPatternHasSlash = MatchName.Contains(TEXT("/"));
	return *this;
}

MCPAssetsBase& MCPAssetsBase::Substring(const FString& InFilter)
{
	MatchName = InFilter;
	bExactMatch = false;
	bPatternHasSlash = MatchName.Contains(TEXT("/"));
	return *this;
}

MCPAssetsBase& MCPAssetsBase::NoDerived()
{
	bNoDerived = true;
	return *this;
}

MCPAssetsBase& MCPAssetsBase::AllContent()
{
	bAllContent = true;
	return *this;
}


bool MCPAssetsBase::Info()
{
	// In theory, there's no reason a person couldn't load/info
	// more than once, to obtain updates.  Might as well allow it.
	AssetResults.Empty();
	UObjectResults.Empty();

	// Query the asset registry
	IAssetRegistry& AR = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry").Get();
	while (AR.IsLoadingAssets()) FPlatformProcess::Sleep(0.1f);

	TArray<FAssetData> Candidates;
	AR.GetAssets(ConfigureFilter(), Candidates);
	for (const FAssetData &Data : Candidates)
	{
		if (AssetMatches(Data)) AssetResults.Add(Data);
		if (bErrorIfAny && (AssetResults.Num() > 0))
		{
			SetError(FString::Printf(TEXT("%s '%s' already exists."), *TargetClass->GetName(), *AssetResults[0].PackageName.ToString()));
			return false;
		}
		if (bErrorIfTwo && (AssetResults.Num() > 1))
		{
			SetError(FString::Printf(
				TEXT("Ambiguous %s name '%s' — matches '%s' and '%s'. Use the full package path to disambiguate."),
				*TargetClass->GetName(), *MatchName, *AssetResults[0].PackageName.ToString(), *AssetResults[1].PackageName.ToString()));
			return false;
		}
		if (AssetResults.Num() >= MaxResults) break;
	}

	// Check error conditions on the result count
	if (bErrorIfNone && AssetResults.IsEmpty())
	{
		SetError(FString::Printf(TEXT("%s '%s' not found."), *TargetClass->GetName(), *MatchName));
		return false;
	}
	return true;
}

bool MCPAssetsBase::Load()
{
	if (!Info()) return false;

	TArray<FAssetData> AssetsToLoad;
	Swap(AssetsToLoad, AssetResults);
	for (const FAssetData &Asset : AssetsToLoad)
	{
		UObject *Obj = TryLoadAsset(Asset);
		if (!Obj) continue;

		// If this is a material open in the editor, use the editor's transient copy.
		if (UMaterial* Mat = Cast<UMaterial>(Obj))
			Obj = MCPUtils::ReplaceMaterialWithTransientCopy(Mat);

		AssetResults.Add(Asset);
		UObjectResults.Add(Obj);
	}
	if (bErrorIfNone && AssetResults.IsEmpty())
	{
		SetError(FString::Printf(TEXT("%s '%s' exists but cannot be loaded."), *TargetClass->GetName(),
					*AssetsToLoad[0].PackageName.ToString()));
		return false;
	}
	return true;
}

FARFilter MCPAssetsBase::ConfigureFilter()
{
	FARFilter Filter;
	for (UClass* C : Scans) Filter.ClassPaths.Add(C->GetClassPathName());
	Filter.bRecursiveClasses = !bNoDerived;
	if (!bAllContent)
	{
		Filter.PackagePaths.Add(FName(TEXT("/Game")));
		Filter.bRecursivePaths = true;
	}
	return Filter;
}

bool MCPAssetsBase::AssetMatches(const FAssetData &Asset)
{
	if (bExactMatch)
	{
		FString Name = bPatternHasSlash ? Asset.PackageName.ToString() : Asset.AssetName.ToString();
		return Name.Equals(MatchName, ESearchCase::IgnoreCase);
	}
	else
	{
		return Asset.AssetName.ToString().Contains(MatchName, ESearchCase::IgnoreCase) ||
			Asset.PackageName.ToString().Contains(MatchName, ESearchCase::IgnoreCase);
	}
}

UObject *MCPAssetsBase::TryLoadAsset(const FAssetData &Asset)
{
	UObject* Obj = Asset.GetAsset();
	if (Obj == nullptr) return nullptr;
	if (Obj->IsA(TargetClass)) return Obj;

	if (TargetClass->IsChildOf(UBlueprint::StaticClass()) &&
		ULevelScriptBlueprint::StaticClass()->IsChildOf(TargetClass))
	{
		UWorld* World = Cast<UWorld>(Obj);
		if (World && World->PersistentLevel)
		{
			ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(true);
			if (LevelBP) return LevelBP;
		}
	}

	return nullptr;
}

void MCPAssetsBase::SetError(const FString &Msg)
{
	AssetResults.Empty();
	UObjectResults.Empty();
	UMCPServer::Printf(TEXT("ERROR: %s\n"), *Msg);
}