terra-llm/Assets/LLM/Editor/GUI/MCPWindow.cs

using UnityEditor;
using UnityEngine;
using LLM.Editor.Core;
using LLM.Editor.Client;
using LLM.Editor.Commands;
using LLM.Editor.Analysis;
using System.Threading.Tasks;
using System.Collections.Generic;

namespace LLM.Editor.GUI
{
    public class MCPWindow : EditorWindow
    {
        private string _promptText = "The grenade launcher should hit the target.";
        private Vector2 _scrollPos;
        private bool _isRequestInProgress;
        
        private readonly List<Object> _stagedContextObjects = new();
        private List<string> _requestedRoles = new();
        private bool _isAnalysisContextRequested;
        private ILlmApiClient _apiClient;

        [MenuItem("Tools/LLM/MCP Assistant")]
        public static void ShowWindow()
        {
            GetWindow<MCPWindow>("MCP Assistant");
        }

        private void OnEnable()
        {
            CommandExecutor.OnQueueUpdated += Repaint;
            RequestAnalysisContextCommand.OnAnalysisContextRequested += HandleAnalysisContextRequested;
            CommandExecutor.OnContextReadyForNextTurn += HandleContextReadyForNextTurn;
            _apiClient = ApiClientFactory.GetClient();
        }

        private void OnDisable()
        {
            CommandExecutor.OnQueueUpdated -= Repaint;
            RequestAnalysisContextCommand.OnAnalysisContextRequested -= HandleAnalysisContextRequested;
            CommandExecutor.OnContextReadyForNextTurn -= HandleContextReadyForNextTurn;
        }

        private void HandleAnalysisContextRequested(RequestAnalysisContextParams contextParams)
        {
            _isAnalysisContextRequested = true;
            _requestedRoles = new List<string>(contextParams.subjectRoles);
            // Ensure our staging list has enough slots for what was requested
            while (_stagedContextObjects.Count < _requestedRoles.Count)
            {
                _stagedContextObjects.Add(null);
            }
            Repaint();
        }
        
        private void HandleContextReadyForNextTurn(string detailedContext)
        {
            Debug.Log("[MCPWindow] Received detailed context. Sending follow-up to LLM.");
            _ = SendFollowUpAsync(detailedContext);
        }

        private void OnGUI()
        {
            EditorGUILayout.LabelField("LLM Co-Pilot", EditorStyles.boldLabel);
            
            DrawSessionManagement();
            EditorGUILayout.Space();

            UnityEngine.GUI.enabled = SessionManager.HasActiveSession();

            
            DrawContextStagingArea();
            EditorGUILayout.Space();

            
            EditorGUILayout.LabelField("Enter your request:");
            _promptText = EditorGUILayout.TextArea(_promptText, GUILayout.Height(60));

            
            DrawActionButton();
            EditorGUILayout.Space();
            
            DrawCommandQueue();
            
            UnityEngine.GUI.enabled = true;
        }
        
        private void DrawSessionManagement()
        {
            EditorGUILayout.BeginHorizontal();
            if (SessionManager.HasActiveSession())
            {
                EditorGUILayout.LabelField($"Session: {SessionManager.GetCurrentSessionId()}");
                if (GUILayout.Button("End Session"))
                {
                    SessionManager.EndSession();
                    _stagedContextObjects.Clear();
                    _isAnalysisContextRequested = false;
                }
            }
            else
            {
                if (GUILayout.Button("Start New Session"))
                {
                    SessionManager.StartNewSession();
                }
            }
            EditorGUILayout.EndHorizontal();
        }
        
        private void DrawContextStagingArea()
        {
            EditorGUILayout.LabelField("Context (Drag & Drop Assets Here)", EditorStyles.boldLabel);
            
            if (GUILayout.Button("Add Context Slot", GUILayout.Width(150)))
            {
                _stagedContextObjects.Add(null);
                if(_isAnalysisContextRequested) _requestedRoles.Add("Custom");
            }

            var indexToRemove = -1; 

            for (var i = 0; i < _stagedContextObjects.Count; i++)
            {
                EditorGUILayout.BeginHorizontal();
                var label = (_isAnalysisContextRequested && i < _requestedRoles.Count) ? _requestedRoles[i] : $"Context {i + 1}";
                _stagedContextObjects[i] = EditorGUILayout.ObjectField(label, _stagedContextObjects[i], typeof(Object), true);
                
                if (GUILayout.Button("X", GUILayout.Width(20)))
                {
                    indexToRemove = i;
                }
                EditorGUILayout.EndHorizontal();
            }

            if (indexToRemove == -1) return;
            _stagedContextObjects.RemoveAt(indexToRemove);
            if(_isAnalysisContextRequested && indexToRemove < _requestedRoles.Count) 
            {
                _requestedRoles.RemoveAt(indexToRemove);
            }
            Repaint();
        }
        
        private void DrawActionButton()
        {
            UnityEngine.GUI.enabled = SessionManager.HasActiveSession() && !_isRequestInProgress;

            if (_isAnalysisContextRequested)
            {
                if (!GUILayout.Button("Provide Staged Context to LLM")) return;
                var contextSummary = ContextBuilder.BuildTier1Summary(_stagedContextObjects);
                var followUpMessage = "Here is the context you requested:\n" + contextSummary;
                _ = SendFollowUpAsync(followUpMessage);
            }
            else
            {
                if (GUILayout.Button(_isRequestInProgress ? "Waiting for response..." : "Send Prompt"))
                {
                    _ = SendInitialPromptAsync();
                }
            }
        }
        
        private void DrawCommandQueue()
        {
            UnityEngine.GUI.enabled = !_isRequestInProgress;
            EditorGUILayout.LabelField("Pending Commands", EditorStyles.boldLabel);
            _scrollPos = EditorGUILayout.BeginScrollView(_scrollPos, EditorStyles.helpBox);
            
            if (CommandExecutor.HasPendingCommands())
            {
                var nextCommand = CommandExecutor.GetNextCommand();
                EditorGUILayout.LabelField("Next Up: " + nextCommand.commandName);

                if (GUILayout.Button("Execute Next Command"))
                {
                    CommandExecutor.ExecuteNextCommand();
                }
            }
            else
            {
                EditorGUILayout.LabelField("No pending commands.");
            }

            EditorGUILayout.EndScrollView();
        }
        
        private async Task SendInitialPromptAsync()
        {
            _isRequestInProgress = true;
            Repaint();
            
            // Pass the staged objects to the API client
            await _apiClient.SendPrompt(_promptText, _stagedContextObjects);

            _isRequestInProgress = false;
            ClearStagingArea();
            Repaint();
        }
        
        private async Task SendFollowUpAsync(string followUpMessage)
        {
            _isRequestInProgress = true;
            Repaint();

            await _apiClient.SendFollowUp(followUpMessage);

            _isRequestInProgress = false;
            
            // Only clear the staging area if the follow-up was initiated by the user
            // providing context. This preserves the UI state otherwise.
            if (_isAnalysisContextRequested)
            {
                ClearStagingArea();
            }
            Repaint();
        }
        
        private void ClearStagingArea()
        {
            _isAnalysisContextRequested = false;
            _stagedContextObjects.Clear();
            _requestedRoles.Clear();
        }
    }
}