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

#include "MCPFetcher.h"
#include "MCPServer.h"
#include "MCPUtils.h"
#include "Engine/Blueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "Engine/SimpleConstructionScript.h"
#include "Engine/SCS_Node.h"
#include "Engine/World.h"
#include "Materials/Material.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "IMaterialEditor.h"
#include "Engine/LevelScriptBlueprint.h"
#include "Subsystems/AssetEditorSubsystem.h"

MCPFetcher::WalkFunc MCPFetcher::GetWalker(const FString& Step)
{
	if (Step.Equals(TEXT("graph"),          ESearchCase::IgnoreCase)) return &MCPFetcher::Graph;
	if (Step.Equals(TEXT("node"),           ESearchCase::IgnoreCase)) return &MCPFetcher::Node;
	if (Step.Equals(TEXT("pin"),            ESearchCase::IgnoreCase)) return &MCPFetcher::Pin;
	if (Step.Equals(TEXT("component"),      ESearchCase::IgnoreCase)) return &MCPFetcher::Component;
	if (Step.Equals(TEXT("levelblueprint"), ESearchCase::IgnoreCase)) return &MCPFetcher::LevelBlueprint;
	return nullptr;
}

void MCPFetcher::PrintDocs()
{
	UMCPServer::Print(TEXT("Some commands take a Path parameter. A Path starts with an asset\n"));
	UMCPServer::Print(TEXT("package path (e.g. /Game/Widgets/WB_Hotkeys), followed by zero or\n"));
	UMCPServer::Print(TEXT("more comma-separated steps that navigate into the asset:\n\n"));
	UMCPServer::Print(TEXT("  graph — Find a named UEdGraph (blank name for material graphs)\n"));
	UMCPServer::Print(TEXT("  node — Find a named UEdGraphNode within a graph or blueprint\n"));
	UMCPServer::Print(TEXT("  pin — Find a named UEdGraphPin on a node\n"));
	UMCPServer::Print(TEXT("  component — Find a named component in a Blueprint's SCS\n"));
	UMCPServer::Print(TEXT("  levelblueprint — Get the level blueprint from a UWorld\n"));
	UMCPServer::Print(TEXT("\nExample: /Game/Widgets/WB_Hotkeys,graph:EventGraph,node:Self_Reference_03,pin:Result\n"));
}


void MCPFetcher::SetObj(UObject* InObj)
{
	UMCPServer::AddTouchedObject(InObj);
	Obj = InObj;
	ResultPin = nullptr;
}

void MCPFetcher::SetPin(UEdGraphPin* InPin)
{
	ResultPin = InPin;
	Obj = nullptr;
}

MCPFetcher& MCPFetcher::SetError()
{
	bError = true;
	Obj = nullptr;
	ResultPin = nullptr;
	OriginalAsset = nullptr;
	Editor = nullptr;
	return *this;
}

MCPFetcher& MCPFetcher::TypeMismatch(const TCHAR* Walker, const TCHAR* Expected)
{
	bError = true;
	if (ResultPin)
		UMCPServer::Printf(TEXT("ERROR: Input to '%s' is a pin, expected %s\n"), Walker, Expected);
	else if (Obj)
		UMCPServer::Printf(TEXT("ERROR: Input to '%s' is %s, expected %s\n"), Walker, *Obj->GetClass()->GetName(), Expected);
	else
		UMCPServer::Printf(TEXT("ERROR: Input to '%s' is null, expected %s\n"), Walker, Expected);
	return *this;
}

MCPFetcher& MCPFetcher::Walk(const FString& Path)
{
	if (bError) return *this;

	TArray<FString> Segments;
	Path.ParseIntoArray(Segments, TEXT(","));
	if (Segments.Num() == 0)
	{
		UMCPServer::Print(TEXT("ERROR: Empty path\n"));
		return SetError();
	}

	for (int32 i = 0; i < Segments.Num(); i++)
	{
		if (!Obj && !ResultPin)
		{
			Asset(Segments[i]);
			if (bError) return *this;
			continue;
		}

		FString Key, Value;
		if (!Segments[i].Split(TEXT(":"), &Key, &Value))
			Key = Segments[i];

		WalkFunc Func = GetWalker(Key);
		if (!Func)
		{
			UMCPServer::Printf(TEXT("ERROR: Unknown path step '%s'\n"), *Key);
			return SetError();
		}
		(this->*Func)(Value);
		if (bError) return *this;
	}

	return *this;
}

MCPFetcher& MCPFetcher::Asset(const FString& PackagePath)
{
	if (bError) return *this;

	SetObj(LoadObject<UObject>(nullptr, *PackagePath));
	if (!Obj)
	{
		UMCPServer::Printf(TEXT("ERROR: Could not load asset '%s'\n"), *PackagePath);
		return SetError();
	}

	OriginalAsset = Obj;

	// Open the editor for this asset (or bring it to front if already open).
	UAssetEditorSubsystem* Sub = GEditor->GetEditorSubsystem<UAssetEditorSubsystem>();
	if (!Sub || !Sub->OpenEditorForAsset(Obj))
	{
		UMCPServer::Printf(TEXT("ERROR: Could not open editor for '%s'\n"), *PackagePath);
		return SetError();
	}
	Editor = Sub->FindEditorForAsset(OriginalAsset, false);
	if (!Editor)
	{
		UMCPServer::Printf(TEXT("ERROR: Could not find editor instance for '%s'\n"), *PackagePath);
		return SetError();
	}

	// If this is a material, use the editor's transient copy.
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		IMaterialEditor *MatEditor = static_cast<IMaterialEditor*>(Editor);
		SetObj(MatEditor->GetMaterialInterface()->GetBaseMaterial());
	}
	return *this;
}

bool MCPFetcher::CheckAssetIsA(UClass* StaticClass)
{
	if (bError) return false;
	if (!OriginalAsset || !OriginalAsset->IsA(StaticClass))
	{
		UMCPServer::Printf(TEXT("ERROR: Asset is %s, expected %s\n"),
			OriginalAsset ? *OriginalAsset->GetClass()->GetName() : TEXT("null"),
			*StaticClass->GetName());
		SetError();
		return false;
	}
	return true;
}

MCPFetcher& MCPFetcher::Graph(const FString& Value)
{
	if (bError) return *this;

	// Material with blank graph name → navigate to the material graph.
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		if (!Value.IsEmpty())
		{
			UMCPServer::Printf(TEXT("ERROR: Materials do not have named graphs (got '%s')\n"), *Value);
			return SetError();
		}
		MCPUtils::EnsureMaterialGraph(Mat);
		if (!Mat->MaterialGraph)
		{
			UMCPServer::Printf(TEXT("ERROR: Material '%s' has no material graph\n"), *Mat->GetName());
			return SetError();
		}
		SetObj(Mat->MaterialGraph);
		return *this;
	}

	UBlueprint* BP = ::Cast<UBlueprint>(Obj);
	if (!BP)
		return TypeMismatch(TEXT("graph"), TEXT("Blueprint or Material"));

	TArray<UEdGraph*> Matches = MCPUtils::AllGraphsNamed(BP, Value);
	if (Matches.Num() == 0)
	{
		UMCPServer::Printf(TEXT("ERROR: Graph '%s' not found in %s\n"), *Value, *BP->GetName());
		return SetError();
	}
	if (Matches.Num() > 1)
	{
		UMCPServer::Printf(TEXT("ERROR: Ambiguous graph '%s' in %s — %d matches\n"), *Value, *BP->GetName(), Matches.Num());
		return SetError();
	}

	SetObj(Matches[0]);
	return *this;
}

MCPFetcher& MCPFetcher::Node(const FString& Value)
{
	if (bError) return *this;

	// If current object is a graph, search that graph
	if (UEdGraph* G = ::Cast<UEdGraph>(Obj))
	{
		UEdGraphNode* Found = nullptr;
		for (UEdGraphNode* N : G->Nodes)
		{
			if (!N || !MCPUtils::Identifies(Value, N))
				continue;
			if (Found)
			{
				UMCPServer::Printf(TEXT("ERROR: Ambiguous node '%s' in graph %s\n"), *Value, *G->GetName());
				return SetError();
			}
			Found = N;
		}
		if (!Found)
		{
			UMCPServer::Printf(TEXT("ERROR: Node '%s' not found in graph %s\n"), *Value, *G->GetName());
			return SetError();
		}
		SetObj(Found);
		return *this;
	}

	// If current object is a blueprint, search all graphs
	if (UBlueprint* BP = ::Cast<UBlueprint>(Obj))
	{
		UEdGraphNode* Found = nullptr;
		for (UEdGraph* G : MCPUtils::AllGraphs(BP))
		{
			for (UEdGraphNode* N : G->Nodes)
			{
				if (!N || !MCPUtils::Identifies(Value, N))
					continue;
				if (Found)
				{
					UMCPServer::Printf(TEXT("ERROR: Ambiguous node '%s' in %s\n"), *Value, *BP->GetName());
					return SetError();
				}
				Found = N;
			}
		}
		if (!Found)
		{
			UMCPServer::Printf(TEXT("ERROR: Node '%s' not found in %s\n"), *Value, *BP->GetName());
			return SetError();
		}
		SetObj(Found);
		return *this;
	}

	return TypeMismatch(TEXT("node"), TEXT("graph or Blueprint"));
}

MCPFetcher& MCPFetcher::Pin(const FString& Value)
{
	if (bError) return *this;

	UEdGraphNode* N = ::Cast<UEdGraphNode>(Obj);
	if (!N)
		return TypeMismatch(TEXT("pin"), TEXT("node"));
	UEdGraphPin* Found = nullptr;
	for (UEdGraphPin *P : N->Pins)
	{
		if (!MCPUtils::Identifies(Value, P))
			continue;
		if (Found)
		{
			UMCPServer::Printf(TEXT("ERROR: Ambiguous pin '%s' on node %s\n"),
					*Value, *MCPUtils::FormatName(N));
			return SetError();
		}
		Found = P;
	}
	if (!Found)
	{
		UMCPServer::Printf(TEXT("ERROR: Pin '%s' not found on node %s\n"),
				*Value, *MCPUtils::FormatName(N));
		return SetError();
	}

	SetPin(Found);
	return *this;
}

MCPFetcher& MCPFetcher::Component(const FString& Value)
{
	if (bError) return *this;

	UBlueprint* BP = ::Cast<UBlueprint>(Obj);
	if (!BP)
		return TypeMismatch(TEXT("component"), TEXT("Blueprint"));

	USimpleConstructionScript* SCS = BP->SimpleConstructionScript;
	if (!SCS)
	{
		UMCPServer::Printf(TEXT("ERROR: Blueprint %s has no SimpleConstructionScript (not an Actor Blueprint)\n"), *BP->GetName());
		return SetError();
	}

	FName SearchName(*Value);
	for (USCS_Node* SCSNode : SCS->GetAllNodes())
	{
		if (SCSNode && SCSNode->GetVariableName() == SearchName)
		{
			SetObj(SCSNode->ComponentTemplate);
			return *this;
		}
	}

	UMCPServer::Printf(TEXT("ERROR: Component '%s' not found in %s\n"), *Value, *BP->GetName());
	return SetError();
}

MCPFetcher& MCPFetcher::LevelBlueprint(const FString& Value)
{
	if (bError) return *this;

	UWorld* World = ::Cast<UWorld>(Obj);
	if (!World)
		return TypeMismatch(TEXT("levelblueprint"), TEXT("World"));

	if (!World->PersistentLevel)
	{
		UMCPServer::Print(TEXT("ERROR: World has no PersistentLevel\n"));
		return SetError();
	}

	ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(true);
	if (!LevelBP)
	{
		UMCPServer::Print(TEXT("ERROR: World has no level blueprint\n"));
		return SetError();
	}

	SetObj(LevelBP);
	return *this;
}

MCPFetcher& MCPFetcher::ToBlueprint()
{
	if (bError) return *this;
	if (::Cast<UBlueprint>(Obj)) return *this;

	if (UWorld* World = ::Cast<UWorld>(Obj))
	{
		if (!World->PersistentLevel)
		{
			UMCPServer::Print(TEXT("ERROR: ToBlueprint: World has no PersistentLevel\n"));
			return SetError();
		}
		ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(true);
		if (!LevelBP)
		{
			UMCPServer::Print(TEXT("ERROR: ToBlueprint: World has no level blueprint\n"));
			return SetError();
		}
		SetObj(LevelBP);
		return *this;
	}

	return TypeMismatch(TEXT("ToBlueprint"), TEXT("Blueprint or World"));
}

MCPFetcher& MCPFetcher::ToGraph()
{
	if (bError) return *this;
	if (::Cast<UEdGraph>(Obj)) return *this;

	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		MCPUtils::EnsureMaterialGraph(Mat);
		if (!Mat->MaterialGraph)
		{
			UMCPServer::Printf(TEXT("ERROR: ToGraph: Material '%s' has no material graph\n"), *Mat->GetName());
			return SetError();
		}
		SetObj(Mat->MaterialGraph);
		return *this;
	}

	return TypeMismatch(TEXT("ToGraph"), TEXT("Graph or Material"));
}