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

#include "MCPServer.h"
#include "MCPHandler.h"
#include "MCPUtils.h"
#include "MCPAssetFinder.h"
#include "UObject/StrongObjectPtr.h"
#include "Materials/MaterialExpression.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "Engine/Blueprint.h"
#include "Engine/World.h"
#include "Engine/Level.h"
#include "Engine/LevelScriptBlueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "EdGraphSchema_K2.h"
#include "K2Node.h"
#include "K2Node_CallFunction.h"
#include "K2Node_Event.h"
#include "K2Node_CustomEvent.h"
#include "K2Node_FunctionEntry.h"
#include "K2Node_EditablePinBase.h"
#include "K2Node_VariableGet.h"
#include "K2Node_VariableSet.h"
#include "K2Node_BreakStruct.h"
#include "K2Node_MakeStruct.h"
#include "K2Node_MacroInstance.h"
#include "K2Node_DynamicCast.h"
#include "K2Node_CallParentFunction.h"
#include "K2Node_IfThenElse.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "Kismet2/KismetEditorUtilities.h"
#include "Dom/JsonValue.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonWriter.h"
#include "Serialization/JsonSerializer.h"
#include "Interfaces/IPv4/IPv4Address.h"
#include "Interfaces/IPv4/IPv4Endpoint.h"
#include "SocketSubsystem.h"
#include "Sockets.h"
#include "Async/Async.h"
#include "UObject/SavePackage.h"
#include "Misc/Paths.h"
#include "Misc/FileHelper.h"
#include "Misc/Guid.h"
#include "AssetToolsModule.h"
#include "IAssetTools.h"
#include "UObject/UObjectIterator.h"
#include "Misc/PackageName.h"
#include "UObject/LinkerLoad.h"
#include "Engine/UserDefinedEnum.h"
#include "Editor.h"
#include "Materials/Material.h"
#include "Materials/MaterialInstanceConstant.h"
#include "Materials/MaterialFunction.h"
#include "Materials/MaterialExpressionScalarParameter.h"
#include "Materials/MaterialExpressionVectorParameter.h"
#include "Materials/MaterialExpressionTextureObjectParameter.h"
#include "Materials/MaterialExpressionTextureSampleParameter2D.h"
#include "Materials/MaterialExpressionStaticSwitchParameter.h"
#include "Materials/MaterialExpressionConstant.h"
#include "Materials/MaterialExpressionConstant2Vector.h"
#include "Materials/MaterialExpressionConstant3Vector.h"
#include "Materials/MaterialExpressionConstant4Vector.h"
#include "Materials/MaterialExpressionTextureSample.h"
#include "Materials/MaterialExpressionTextureCoordinate.h"
#include "Materials/MaterialExpressionComponentMask.h"
#include "Materials/MaterialExpressionCustom.h"
#include "Materials/MaterialExpressionAppendVector.h"
#include "Materials/MaterialExpressionAdd.h"
#include "Materials/MaterialExpressionMultiply.h"
#include "Materials/MaterialExpressionLinearInterpolate.h"
#include "Materials/MaterialExpressionClamp.h"
#include "Materials/MaterialExpressionOneMinus.h"
#include "Materials/MaterialExpressionPower.h"
#include "Materials/MaterialExpressionTime.h"
#include "Materials/MaterialExpressionWorldPosition.h"
#include "Materials/MaterialExpressionFunctionInput.h"
#include "Materials/MaterialExpressionFunctionOutput.h"
#include "Materials/MaterialExpressionMaterialFunctionCall.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "MaterialGraph/MaterialGraphSchema.h"

// Animation Blueprint support
#include "Animation/AnimBlueprint.h"
#include "Animation/AnimBlueprintGeneratedClass.h"
#include "Animation/Skeleton.h"
#include "AnimGraphNode_StateMachine.h"
#include "AnimGraphNode_AssetPlayerBase.h"
#include "AnimGraphNode_SequencePlayer.h"
#include "AnimGraphNode_BlendSpacePlayer.h"
#include "AnimGraphNode_Base.h"
#include "AnimStateNode.h"
#include "AnimStateTransitionNode.h"
#include "AnimStateConduitNode.h"
#include "AnimStateEntryNode.h"
#include "AnimationStateMachineGraph.h"
#include "AnimationGraph.h"
#include "AnimationTransitionGraph.h"

// ============================================================
// Get() — retrieve the active server via the editor subsystem
// ============================================================

#include "MCPEditorSubsystem.h"

FMCPServer* FMCPServer::Get()
{
	if (!GEditor) return nullptr;
	auto* Sub = GEditor->GetEditorSubsystem<UBlueprintMCPEditorSubsystem>();
	return Sub ? Sub->GetServer() : nullptr;
}


// ============================================================
// SEH wrappers for crash-safe compilation and saving.
// MSVC constraint: __try/__except functions must NOT contain C++
// objects with destructors. We factor the actual work into
// separate "inner" functions and only do try/except in thin wrappers.
// ============================================================
#if PLATFORM_WINDOWS

// Inner functions that do the actual C++ work (may have destructors)
static void CompileBlueprintInner(UBlueprint* BP, EBlueprintCompileOptions Opts)
{
	FKismetEditorUtilities::CompileBlueprint(BP, Opts, nullptr);
}

static ESavePackageResult SavePackageInner(
	UPackage* Package, UObject* Asset, const TCHAR* Filename,
	FSavePackageArgs* SaveArgs)
{
	FSavePackageResultStruct Result = UPackage::Save(Package, Asset, Filename, *SaveArgs);
	return Result.Result;
}

// SEH wrappers — absolutely NO C++ objects with destructors here.
// EXCEPTION_EXECUTE_HANDLER = 1 (avoiding Windows.h include)
#pragma warning(push)
#pragma warning(disable: 4611) // interaction between '_setjmp' and C++ object destruction
int32 TryCompileBlueprintSEH(UBlueprint* BP, EBlueprintCompileOptions Opts)
{
	__try
	{
		CompileBlueprintInner(BP, Opts);
		return 0;
	}
	__except (1)
	{
		return -1;
	}
}

static int32 TrySavePackageSEH(
	UPackage* Package, UObject* Asset, const TCHAR* Filename,
	FSavePackageArgs* SaveArgs, ESavePackageResult* OutResult)
{
	__try
	{
		*OutResult = SavePackageInner(Package, Asset, Filename, SaveArgs);
		return 0;
	}
	__except (1)
	{
		*OutResult = ESavePackageResult::Error;
		return -1;
	}
}


// Inner: create expression, register in material, and trigger PostEditChange.
// All of this may crash for classes that are effectively abstract.
static void AddMaterialExpressionInner(
	UObject* Owner, UClass* ExprClass, UMaterial* Material, UMaterialFunction* MatFunc,
	int32 PosX, int32 PosY, UMaterialExpression** OutExpr)
{
	*OutExpr = NewObject<UMaterialExpression>(Owner, ExprClass);
	if (!*OutExpr) return;

	(*OutExpr)->MaterialExpressionEditorX = PosX;
	(*OutExpr)->MaterialExpressionEditorY = PosY;

	if (Material)
	{
		Material->GetExpressionCollection().AddExpression(*OutExpr);
		if (Material->MaterialGraph)
		{
			Material->MaterialGraph->RebuildGraph();
		}
		Material->PreEditChange(nullptr);
		Material->PostEditChange();
		Material->MarkPackageDirty();
	}
	else if (MatFunc)
	{
		MatFunc->GetExpressionCollection().AddExpression(*OutExpr);
		MatFunc->PreEditChange(nullptr);
		MatFunc->PostEditChange();
		MatFunc->MarkPackageDirty();
	}
}

// Inner: remove a bad expression from a material after a crash
static void CleanupBadExpressionInner(UMaterial* Material, UMaterialFunction* MatFunc, UMaterialExpression* BadExpr)
{
	if (!BadExpr) return;
	if (Material)
	{
		Material->GetExpressionCollection().RemoveExpression(BadExpr);
		if (Material->MaterialGraph)
		{
			Material->MaterialGraph->RebuildGraph();
		}
	}
	else if (MatFunc)
	{
		MatFunc->GetExpressionCollection().RemoveExpression(BadExpr);
	}
	BadExpr->MarkAsGarbage();
}

int32 TryAddMaterialExpressionSEH(
	UObject* Owner, UClass* ExprClass, UMaterial* Material, UMaterialFunction* MatFunc,
	int32 PosX, int32 PosY, UMaterialExpression** OutExpr)
{
	__try
	{
		AddMaterialExpressionInner(Owner, ExprClass, Material, MatFunc, PosX, PosY, OutExpr);
		return 0;
	}
	__except (1)
	{
		// Try to clean up the partially-added expression
		__try
		{
			CleanupBadExpressionInner(Material, MatFunc, *OutExpr);
		}
		__except (1)
		{
			// Cleanup also crashed — nothing more we can do
		}
		*OutExpr = nullptr;
		return -1;
	}
}
#pragma warning(pop)

#endif // PLATFORM_WINDOWS


void FMCPServer::DispatchToolCall(const FString& ToolName, const FJsonObject* Params, FJsonObject* Result)
{
	UMCPAssetFinder::Refresh();
	if (UClass** HandlerClass = MCPHandlerRegistry.Find(ToolName))
	{
		const bool bIsMutation = MutationEndpoints.Contains(ToolName);
		if (bIsMutation && GEditor)
		{
			GEditor->BeginTransaction(FText::FromString(FString::Printf(TEXT("BlueprintMCP: %s"), *ToolName)));
		}

		TStrongObjectPtr<UObject> HandlerObj(NewObject<UObject>(GetTransientPackage(), *HandlerClass));
		IMCPHandler* Handler = Cast<IMCPHandler>(HandlerObj.Get());
		FString PopulateError = MCPUtils::PopulateFromJson(HandlerObj->GetClass(), HandlerObj.Get(), Params);
		if (PopulateError.IsEmpty())
		{
			Handler->Handle(Params, Result);
		}
		else
		{
			MCPUtils::MakeErrorJson(Result, PopulateError);
		}

		if (bIsMutation && GEditor)
		{
			GEditor->EndTransaction();
		}
	}
	else
	{
		MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Unknown tool: %s"), *ToolName));
	}
}

// ============================================================
// HandleRequest — parse a JSON command and dispatch
// ============================================================

FString FMCPServer::HandleRequest(const FString& Line)
{
	TSharedPtr<FJsonObject> Request = MCPUtils::ParseBodyJson(Line);
	if (!Request.IsValid())
	{
		TSharedRef<FJsonObject> ErrResult = MakeShared<FJsonObject>();
		MCPUtils::MakeErrorJson(&*ErrResult, TEXT("JSON parse error"));
		return MCPUtils::JsonToString(ErrResult);
	}

	FString Command;
	if (!Request->TryGetStringField(TEXT("command"), Command))
	{
		TSharedRef<FJsonObject> ErrResult = MakeShared<FJsonObject>();
		MCPUtils::MakeErrorJson(&*ErrResult, TEXT("Missing 'command' field"));
		return MCPUtils::JsonToString(ErrResult);
	}
	Request->RemoveField(TEXT("command"));

	TSharedRef<FJsonObject> Result = MakeShared<FJsonObject>();
	DispatchToolCall(Command, Request.Get(), &*Result);
	return MCPUtils::JsonToString(Result);
}

// ============================================================
// TCP Server: Start / Stop / ProcessOneRequest
// ============================================================

bool FMCPServer::Start(int32 InPort, bool bEditorMode)
{
	Port = InPort;
	bIsEditor = bEditorMode;

	// Register handlers
	BuildMCPHandlerRegistry();
	RegisterHandlers();

	// Create TCP listen socket
	ISocketSubsystem* SocketSub = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);
	ListenSocket = SocketSub->CreateSocket(NAME_Stream, TEXT("MCPServer"), false);
	if (!ListenSocket)
	{
		UE_LOG(LogTemp, Error, TEXT("BlueprintMCP: Failed to create listen socket"));
		return false;
	}

	ListenSocket->SetReuseAddr(true);
	ListenSocket->SetNonBlocking(true);

	TSharedRef<FInternetAddr> Addr = SocketSub->CreateInternetAddr();
	bool bIsValid = false;
	Addr->SetIp(TEXT("127.0.0.1"), bIsValid);
	Addr->SetPort(Port);

	if (!ListenSocket->Bind(*Addr))
	{
		UE_LOG(LogTemp, Error, TEXT("BlueprintMCP: Failed to bind to port %d"), Port);
		SocketSub->DestroySocket(ListenSocket);
		ListenSocket = nullptr;
		return false;
	}

	if (!ListenSocket->Listen(4))
	{
		UE_LOG(LogTemp, Error, TEXT("BlueprintMCP: Failed to listen on port %d"), Port);
		SocketSub->DestroySocket(ListenSocket);
		ListenSocket = nullptr;
		return false;
	}

	bRunning = true;
	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: MCP server listening on tcp://localhost:%d"), Port);
	return true;
}

void FMCPServer::Stop()
{
	if (!bRunning) return;

	ISocketSubsystem* SocketSub = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);

	// Set shutdown flag and drain pending messages under lock
	{
		FScopeLock Lock(&Mutex);
		bShuttingDown = true;
		for (auto& Msg : PendingMessages)
		{
			Msg->Response.SetValue(FString());
		}
		PendingMessages.Empty();
	}

	// Close all client sockets (unblocks their blocking reads)
	for (auto& Client : Clients)
	{
		if (Client->Socket)
		{
			Client->Socket->Close();
		}
	}

	// Wait for client threads to exit
	for (auto& Client : Clients)
	{
		Client->ThreadFuture.Wait();
		if (Client->Socket)
		{
			SocketSub->DestroySocket(Client->Socket);
		}
	}
	Clients.Empty();

	// Close listen socket
	if (ListenSocket)
	{
		ListenSocket->Close();
		SocketSub->DestroySocket(ListenSocket);
		ListenSocket = nullptr;
	}

	bRunning = false;
	bShuttingDown = false;
	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Server stopped."));
}

void FMCPServer::AcceptNewConnections()
{
	if (!ListenSocket) return;

	bool bHasPending = false;
	if (!ListenSocket->HasPendingConnection(bHasPending) || !bHasPending) return;

	FSocket* ClientSocket = ListenSocket->Accept(TEXT("MCPClient"));
	if (!ClientSocket) return;

	ClientSocket->SetNonBlocking(false); // client threads use blocking I/O

	TSharedPtr<FClientConnection> Client = MakeShared<FClientConnection>();
	Client->Socket = ClientSocket;
	Client->ThreadFuture = Async(EAsyncExecution::Thread, [this, Client]() { ClientThreadFunc(this, Client); });
	Clients.Add(Client);

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Client connected."));
}

void FMCPServer::CleanupFinishedClients()
{
	ISocketSubsystem* SocketSub = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);

	for (int32 i = Clients.Num() - 1; i >= 0; --i)
	{
		if (!Clients[i]->bDone) continue;

		Clients[i]->ThreadFuture.Wait();
		if (Clients[i]->Socket)
		{
			SocketSub->DestroySocket(Clients[i]->Socket);
		}
		Clients.RemoveAt(i);
	}
}

void FMCPServer::ClientThreadFunc(FMCPServer* Server, TSharedPtr<FClientConnection> Client)
{
	FSocket* Socket = Client->Socket;
	FString LineBuffer;
	uint8 RecvBuf[4096];

	while (true)
	{
		int32 BytesRead = 0;
		if (!Socket->Recv(RecvBuf, sizeof(RecvBuf) - 1, BytesRead))
		{
			break; // socket error or closed
		}
		if (BytesRead <= 0)
		{
			break; // connection closed
		}

		RecvBuf[BytesRead] = 0;
		LineBuffer += UTF8_TO_TCHAR((const ANSICHAR*)RecvBuf);

		// Process complete lines
		int32 NewlineIdx;
		while (LineBuffer.FindChar(TEXT('\n'), NewlineIdx))
		{
			FString Line = LineBuffer.Left(NewlineIdx).TrimEnd();
			LineBuffer.RightChopInline(NewlineIdx + 1);

			if (Line.IsEmpty()) continue;

			// Enqueue the line for game-thread processing
			TSharedPtr<FMCPServer::FPendingMessage> Msg = MakeShared<FMCPServer::FPendingMessage>();
			Msg->Line = Line;
			TFuture<FString> Future = Msg->Response.GetFuture();

			{
				FScopeLock Lock(&Server->Mutex);
				if (Server->bShuttingDown)
				{
					Client->bDone = true;
					return;
				}
				Server->PendingMessages.Add(Msg);
			}

			// Block until the game thread processes this message
			FString Response = Future.Get();

			// Empty response means either notification or shutdown
			if (Response.IsEmpty()) continue;

			// Write the response line back (blocking)
			FString ResponseLine = Response + TEXT("\n");
			FTCHARToUTF8 Utf8(*ResponseLine);
			int32 BytesSent = 0;
			Socket->Send((const uint8*)Utf8.Get(), Utf8.Length(), BytesSent);
		}
	}

	Client->bDone = true;
	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Client disconnected."));
}

bool FMCPServer::ProcessOneRequest()
{
	// Accept new connections (non-blocking)
	AcceptNewConnections();

	// Clean up finished client threads
	CleanupFinishedClients();

	// Dequeue one pending message
	TSharedPtr<FPendingMessage> Msg;
	{
		FScopeLock Lock(&Mutex);
		if (PendingMessages.Num() > 0)
		{
			Msg = PendingMessages[0];
			PendingMessages.RemoveAt(0);
		}
	}

	if (!Msg.IsValid()) return false;

	// Process on game thread
	FString Response = HandleRequest(Msg->Line);
	Msg->Response.SetValue(Response);

	return true;
}

// ============================================================
// RegisterHandlers / BuildMCPHandlerRegistry
// ============================================================

void FMCPServer::RegisterHandlers()
{
	// Mutation endpoints — wrapped in undo transactions by DispatchToolCall()
	MutationEndpoints = {
		TEXT("replace_function_calls_in_blueprint"),
		TEXT("change_blueprint_variable_type"),
		TEXT("change_function_parameter_type"),
		TEXT("remove_function_parameter"),
		TEXT("delete_asset"),
		TEXT("connect_blueprint_pins"),
		TEXT("disconnect_blueprint_pins"),
		TEXT("refresh_all_nodes_in_graph"),
		TEXT("set_pin_default_values"),
		TEXT("set_node_positions"),
		TEXT("change_struct_node_type"),
		TEXT("delete_node_from_graph"),
		TEXT("duplicate_nodes_in_graph"),
		TEXT("spawn_nodes_in_graph"),
		TEXT("set_node_comment"),
		TEXT("rename_asset"),
		TEXT("reparent_blueprint"),
		TEXT("set_class_default_value"),
		TEXT("create_blueprint_asset"),
		TEXT("create_blueprint_graph"),
		TEXT("delete_blueprint_graph"),
		TEXT("rename_blueprint_graph"),
		TEXT("add_blueprint_variable"),
		TEXT("remove_blueprint_variable"),
		TEXT("set_blueprint_variable_metadata"),
		TEXT("add_blueprint_interface"),
		TEXT("remove_blueprint_interface"),
		TEXT("add_event_dispatcher"),
		TEXT("add_function_parameter"),
		TEXT("add_blueprint_component"),
		TEXT("remove_blueprint_component"),
		TEXT("create_material_asset"),
		TEXT("set_material_property"),
		TEXT("add_material_expression"),
		TEXT("delete_material_expression"),
		TEXT("connect_material_expression_pins"),
		TEXT("disconnect_material_expression_pin"),
		TEXT("set_material_expression_property"),
		TEXT("set_material_expression_position"),
		TEXT("create_material_instance_asset"),
		TEXT("set_material_instance_parameter"),
		TEXT("reparent_material_instance"),
		TEXT("create_material_function_asset"),
		TEXT("add_anim_state_to_machine"),
		TEXT("remove_anim_state_from_machine"),
		TEXT("add_anim_state_transition"),
		TEXT("set_anim_transition_rule"),
		TEXT("set_anim_state_animation"),
	};

}

void FMCPServer::BuildMCPHandlerRegistry()
{
	for (TObjectIterator<UClass> It; It; ++It)
	{
		UClass* Class = *It;
		if (!Class->ImplementsInterface(UMCPHandler::StaticClass()))
		{
			continue;
		}
		if (Class->HasAnyClassFlags(CLASS_Abstract))
		{
			continue;
		}
		const FString& ToolName = Class->GetMetaData(TEXT("ToolName"));
		if (ToolName.IsEmpty())
		{
			UE_LOG(LogTemp, Warning, TEXT("BlueprintMCP: %s has no ToolName meta — skipping."), *Class->GetName());
			continue;
		}
		if (MCPHandlerRegistry.Contains(ToolName))
		{
			UE_LOG(LogTemp, Warning, TEXT("BlueprintMCP: Duplicate ToolName '%s' on %s — skipping."), *ToolName, *Class->GetName());
			continue;
		}
		MCPHandlerRegistry.Add(ToolName, Class);
		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Registered handler '%s' → %s"), *ToolName, *Class->GetName());
	}

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: %d new-style handlers registered."), MCPHandlerRegistry.Num());
}