integration/Plugins/UEWingman/Source/UEWingman/Private/WingServer.cpp

#include "WingServer.h"
#include "WingHandler.h"
#include "WingProperty.h"
#include "WingManual.h"
#include "WingLogCapture.h"
#include "WingUtils.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"

UWingServer* UWingServer::GWingServer = nullptr;

// ============================================================
// Initialization and Shutdown
// ============================================================

void UWingServer::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);
	GWingServer = this;

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

	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("UEWingman: Failed to bind to port %d"), Port);
		SocketSub->DestroySocket(ListenSocket);
		ListenSocket = nullptr;
		return;
	}

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

	BuildWingHandlerRegistry();
	LogCapture.bEnabled = false;
	LogCapture.Install();
	bRunning = true;
	UE_LOG(LogTemp, Display, TEXT("UEWingman: MCP server listening on tcp://localhost:%d"), Port);
}

void UWingServer::Deinitialize()
{
	if (!bRunning)
	{
		Super::Deinitialize();
		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;
	}

	LogCapture.Uninstall();
	bRunning = false;
	bShuttingDown = false;
	GWingServer = nullptr;
	UE_LOG(LogTemp, Display, TEXT("UEWingman: Server stopped."));
	Super::Deinitialize();
}

// ============================================================
// FTickableEditorObject interface
// ============================================================

void UWingServer::Tick(float DeltaTime)
{
	// Accept new connections (non-blocking)
	AcceptNewConnections();

	// Clean up finished client threads
	CleanupFinishedClients();

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

	// If we have a request, process it.
	if (Request.IsValid())
	{
		FString Response = HandleRequest(Request->Line);
		Request->Response.SetValue(Response);
	}
}

void UWingServer::TickServer(float DeltaTime)
{
	if (GWingServer) GWingServer->Tick(DeltaTime);
}

bool UWingServer::IsTickable() const
{
	return bRunning;
}

TStatId UWingServer::GetStatId() const
{
	RETURN_QUICK_DECLARE_CYCLE_STAT(UWingServer, STATGROUP_Tickables);
}

// ============================================================
// HandleRequest — Given a command, execute it.
// ============================================================

FString UWingServer::HandleRequest(const FString& Line)
{
	LogCapture.CapturedErrors.Empty();
	LogCapture.bEnabled = true;
	HandlerOutput.Reset();
	SuggestedManualSections.Empty();
	LastHandler = nullptr;

	TryCallHandler(Line);

	Notifier.SendNotifications();
	LogCapture.bEnabled = false;
	for (const FString& Msg : LogCapture.CapturedErrors)
	{
		UWingServer::Printf(TEXT("UE_LOG: %s\n"), *Msg);
	}
	LogCapture.CapturedErrors.Empty();
	if (!SuggestedManualSections.IsEmpty())
	{
		UWingServer::SuggestManual(WingManual::Section::HandlerHelp);
		WingManual::PrintManual(SuggestedManualSections, LastHandler, true);
	}
	FString Result = HandlerOutput.ToString();
	HandlerOutput.Reset();
	for (int32 i = 0; i < Result.Len(); ++i)
	{
		if (Result[i] == TEXT('\0')) Result[i] = TEXT(' ');
	}
	return Result;
}

void UWingServer::TryCallHandler(const FString &Line)
{
	// Turn the request string into a JSON tree.
	TSharedPtr<FJsonObject> Request;
	TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Line);
	FJsonSerializer::Deserialize(Reader, Request);
	if (!Request.IsValid())
	{
		UWingServer::Printf(TEXT("Request is not valid JSON"));
		return;
	}

	// Extract the command from the request.
	FString Command;
	if (!Request->TryGetStringField(TEXT("command"), Command))
	{
		UWingServer::Printf(TEXT("Request does not contain 'command' parameter"));
		UWingServer::Printf(TEXT("We recommend sending command='UserManual'."));
		return;
	}
	Request->RemoveField(TEXT("command"));

	// Find the handler for the specified command.
	FWingHandlerConfig* Found = FindHandler(Command);
	if (!Found)
	{
		UWingServer::Printf(TEXT("Unknown command: %s"), *Command);
		UWingServer::SuggestManual(WingManual::Section::ImportantCommands);
		return;
	}
	LastHandler = Found;

	// Make an object of the handler class.
	TStrongObjectPtr<UObject> HandlerObj(NewObject<UObject>(GetTransientPackage(), Found->Class.Get()));
	UWingHandler* Handler = Cast<UWingHandler>(HandlerObj.Get());
	Handler->ConfigurationObject = Found->Config.Get();

	// Populate the handler object with the request parameters.
	if (!FWingProperty::PopulateFromJson(HandlerObj->GetClass(), HandlerObj.Get(), &*Request))
	{
		UWingServer::SuggestManual(WingManual::Section::HandlerHelp);
		return;
	}

	// MCP handlers must not run inside an undo transaction.
	check(GUndo == nullptr);

	// Invoke the handler.
	Handler->Handle();
}

// ============================================================
// Connection Maintenance
// ============================================================

void UWingServer::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("UEWingman: Client connected."));
}

void UWingServer::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 UWingServer::ClientThreadFunc(UWingServer* 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;

			// Wait for the asset registry to finish its initial scan.
			{
				IAssetRegistry& AR = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry").Get();
				while (AR.IsLoadingAssets()) FPlatformProcess::Sleep(0.25f);
			}

			// Enqueue the line for game-thread processing
			TSharedPtr<UWingServer::FPendingMessage> Msg = MakeShared<UWingServer::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();

			// Write the response back, null-terminated (blocking)
			FTCHARToUTF8 Utf8(*Response);
			int32 BytesSent = 0;
			Socket->Send((const uint8*)Utf8.Get(), Utf8.Length() + 1, BytesSent);
		}
	}

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


// ============================================================
// BuildWingHandlerRegistry
// ============================================================

void UWingServer::AddHandler(UObject* Obj, const FString& Documentation)
{
	AddHandler(Obj, WingUtils::GetHandlerName(Obj->GetClass()), nullptr, EWingHandlerKind::Normal, Documentation);
}

void UWingServer::AddHandler(UObject* Obj, const FString& Name, UObject* Config, EWingHandlerKind Kind, const FString& Documentation)
{
	FWingHandlerConfig H;
	H.Name = Name;
	H.Documentation = Documentation;
	H.Class = TStrongObjectPtr<UClass>(Obj->GetClass());
	H.Config = TStrongObjectPtr<UObject>(Config);
	H.Kind = Kind;

	GWingServer->WingHandlerRegistry.Add(MoveTemp(H));
}

void UWingServer::BuildWingHandlerRegistry()
{
	for (UClass* Class : WingUtils::CollectHandlerClasses())
	{
		UWingHandler* CDO = Cast<UWingHandler>(Class->GetDefaultObject());
		CDO->Register();
	}
	WingHandlerRegistry.Sort([](const FWingHandlerConfig& A, const FWingHandlerConfig& B) { return A.Name < B.Name; });
}

FWingHandlerConfig* UWingServer::FindHandler(const FString& Name)
{
	int32 Index = Algo::LowerBoundBy(WingHandlerRegistry, Name, [](const FWingHandlerConfig& H) { return H.Name; });
	if (Index < WingHandlerRegistry.Num() && WingHandlerRegistry[Index].Name == Name)
	{
		return &WingHandlerRegistry[Index];
	}
	return nullptr;
}