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

#include "WingServer.h"
#include "WingProperty.h"
#include "WingUtils.h"
#include "UObject/StrongObjectPtr.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "Misc/CoreDelegates.h"
#include "Misc/OutputDeviceRedirector.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"
#include "SocketSubsystem.h"
#include "Sockets.h"
#include "Async/Async.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;
	}

	LoadingPhasesCompleteHandle = FCoreDelegates::OnAllModuleLoadingPhasesComplete.AddUObject(this, &UWingServer::BuildWingHandlerRegistry);
	LogCapture.bEnabled = false;
	GLog->AddOutputDevice(&LogCapture);
	bRunning = true;
	UE_LOG(LogTemp, Display, TEXT("UEWingman: MCP server listening on tcp://localhost:%d"), Port);
}

void UWingServer::Deinitialize()
{
	FCoreDelegates::OnAllModuleLoadingPhasesComplete.Remove(LoadingPhasesCompleteHandle);

	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(TArray<uint8>());
		}
		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;
	}

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

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

void UWingServer::Tick(float DeltaTime)
{
	if (!bRunning) return;

	// 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())
	{
		TArray<uint8> Response = HandleRequest(Request->Request);
		Request->Response.SetValue(Response);
	}
}

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

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

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

TArray<uint8> UWingServer::HandleRequest(const TArray<uint8>& RequestBytes)
{
	TArray<FString> Argv;
	FString ResponseText;

	if (DeserializeArgv(RequestBytes, Argv))
	{
		PreCallHandler();
		TryCallHandler(Argv);
		ResponseText = PostCallHandler();
	}
	else ResponseText = TEXT("Invalid argv encoding (bug in ue-wingman.py)\n");

	FTCHARToUTF8 Utf8(*ResponseText);
	return TArray<uint8>(reinterpret_cast<const uint8*>(Utf8.Get()), Utf8.Length());
}

void UWingServer::PreCallHandler()
{
	LogCapture.CapturedErrors.Empty();
	LogCapture.bEnabled = true;
	WingOut::StdoutBuffer.Reset();
	SuggestedManualSections.Empty();
	bSuggestHandlerHelp = false;
	LastHandler = nullptr;
}

FString UWingServer::PostCallHandler()
{
	Notifier.SendNotifications();
	LogCapture.bEnabled = false;
	for (const FString& Msg : LogCapture.CapturedErrors)
	{
		WingOut::Stdout.Printf(TEXT("UE_LOG: %s\n"), *Msg);
	}
	LogCapture.CapturedErrors.Empty();
	if (bSuggestHandlerHelp || (!SuggestedManualSections.IsEmpty()))
	{
		if (LastHandler) WingManual::PrintHandlerHelp(*LastHandler);
		if ((LastHandler == nullptr) || (LastHandler->Name != TEXT("Documentation_Manual")))
		{
			WingOut::Stdout.Print(TEXT("To see manual: command=Documentation_Manual\n"));
		}
		if (!SuggestedManualSections.IsEmpty())
		{
			WingManual::PrintSectionNames(TEXT("Suggested manual sections: "),
				SuggestedManualSections, WingOut::Stdout);
		}
	}
	FString Result = WingOut::StdoutBuffer.ToString();
	WingOut::StdoutBuffer.Reset();
	return Result;
}

void UWingServer::TryCallHandler(TArrayView<const FString> Argv)
{
	FString Command = "Documentation_Manual";
	if (Argv.Num() > 0)
	{
		Command = Argv[0];
		Argv = Argv.RightChop(1);
	}

	if ((Command.Equals(TEXT("--help"))) ||
		(Command.Equals(TEXT("-help"))) ||
		(Command.Equals(TEXT("help"))))
	{
		Command = "Documentation_Manual";
    }

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

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

	// Populate the handler object with argv parameters.
	TArray<FWingProperty> Props = FWingProperty::GetVisible(Handler, true);
	if (!FWingProperty::PopulateFromArgv(Props, Argv, WingOut::Stdout))
	{
		UWingServer::SuggestHandlerHelp();
		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);
}

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);
	}
}

// ============================================================
// Stuff Performed on the Client Thread
// ============================================================

void UWingServer::ClientThreadFunc(UWingServer* Server, TSharedPtr<FClientConnection> Client)
{
	FSocket* Socket = Client->Socket;

	WaitForAssetRegistry();

	TArray<uint8> Request;
	if (!ReceiveRequest(Socket, Request))
	{
		Client->bDone = true;
		return;
	}

	TArray<uint8> Response;
	if (!ProcessRequestOnGameThread(Request, Response))
	{
		Client->bDone = true;
		return;
	}

	SendAll(Socket, Response.GetData(), Response.Num());
	Client->bDone = true;
}

uint32 UWingServer::UnpackBigEndian(const uint8 *Data)
{
	return 
		((uint32)Data[0] << 24) |
		((uint32)Data[1] << 16) |
		((uint32)Data[2] << 8) |
		(uint32)Data[3];
}

bool UWingServer::DeserializeArgv(
	const TArray<uint8>& RequestBytes, TArray<FString>& Argv)
{
	Argv.Empty();

	int32 Offset = 0;
	while (Offset < RequestBytes.Num())
	{
		if (RequestBytes.Num() - Offset < 4)
		{
			Argv.Empty();
			return false;
		}

		uint32 Length = UnpackBigEndian(RequestBytes.GetData() + Offset);
		Offset += 4;

		if ((uint32)(RequestBytes.Num() - Offset) < Length)
		{
			Argv.Empty();
			return false;
		}

		Argv.Add(FString::ConstructFromPtrSize(
			reinterpret_cast<const UTF8CHAR*>(RequestBytes.GetData() + Offset),
			Length));
		Offset += (int32)Length;
	}

	return true;
}

bool UWingServer::ReceiveRequest(FSocket* Socket, TArray<uint8>& OutRequest)
{
	constexpr int32 MaxRecvBufBytes = 1024 * 1024;
	constexpr int32 ChunkSize = 8192;

	TArray<uint8> RecvBuf;
	RecvBuf.Reserve(ChunkSize);

	// Unreal's FSocket API is fundamentally broken: recv cannot
	// differentiate between a socket that has been cleanly closed
	// and a socket that has had an error.  So we have no choice
	// but to just read until recv returns false (which could be a
	// clean close or an error).  Then, we check if we have a cleanly
	// encoded payload: if so, we assume everything is fine.
	while (true)
	{
		uint8 Temp[ChunkSize];
		int32 BytesRead = 0;
		if (!Socket->Recv(Temp, ChunkSize, BytesRead))
		{
			break;
		}
		if (BytesRead <= 0) break;
		if (RecvBuf.Num() + BytesRead > MaxRecvBufBytes)
		{
			return false;
		}
		RecvBuf.Append(Temp, BytesRead);
	}

	if (RecvBuf.Num() < 4) return false;
	uint32 Size = UnpackBigEndian(RecvBuf.GetData());
	if ((uint32)RecvBuf.Num() != (4u + Size)) return false;
	RecvBuf.RemoveAt(0, 4);

	OutRequest = MoveTemp(RecvBuf);
	return true;
}

bool UWingServer::SendAll(FSocket* Socket, const uint8* Data, int32 BytesToSend)
{
	while (BytesToSend > 0)
	{
		int32 BytesSent = 0;
		if (!Socket->Send(Data, BytesToSend, BytesSent) || (BytesSent <= 0))
		{
			return false;
		}
		Data += BytesSent;
		BytesToSend -= BytesSent;
	}
	return true;
}

bool UWingServer::ProcessRequestOnGameThread(
	const TArray<uint8>& Request, TArray<uint8>& Response)
{
	// Enqueue the message for game-thread processing.
	TSharedPtr<UWingServer::FPendingMessage> Msg =
		MakeShared<UWingServer::FPendingMessage>();
	Msg->Request = Request;
	TFuture<TArray<uint8>> Future = Msg->Response.GetFuture();

	{
		FScopeLock Lock(&GWingServer->Mutex);
		if (GWingServer->bShuttingDown)
		{
			return false;
		}
		GWingServer->PendingMessages.Add(Msg);
	}

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

void UWingServer::WaitForAssetRegistry()
{
	IAssetRegistry& AR =
		FModuleManager::LoadModuleChecked<FAssetRegistryModule>(
			"AssetRegistry").Get();
	while (AR.IsLoadingAssets()) FPlatformProcess::Sleep(0.25f);
}

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

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

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

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

void UWingServer::BuildWingHandlerRegistry()
{
	WingHandlerRegistry.Empty();
	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;
}

TStringBuilder<65536> WingOut::StdoutBuffer;
WingOut WingOut::Stdout(&WingOut::StdoutBuffer);
WingOut WingOut::None(nullptr);