using System.Linq; using UnityEditor; using UnityEngine; using LLM.Editor.Core; using System.Threading; using LLM.Editor.Client; using LLM.Editor.Helper; 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 _stagedContextObjects = new(); private List _requestedRoles = new(); private bool _isAnalysisContextRequested; private ILlmApiClient _apiClient; private CancellationTokenSource _cancellationTokenSource; private List _currentPlan = new(); // --- Fields for Dummy API Client UI --- private List _dummyTriggerPhrases = new(); private int _selectedDummyTriggerIndex; private bool _dummyResponsesLoaded; [MenuItem("LLM/MCP Assistant")] public static void ShowWindow() { GetWindow("MCP Assistant"); } private void OnEnable() { CommandExecutor.OnQueueUpdated += Repaint; CommandExecutor.OnPlanUpdated += HandlePlanUpdated; RequestAnalysisContextCommand.OnAnalysisContextRequested += HandleAnalysisContextRequested; CommandExecutor.OnContextReadyForNextTurn += HandleContextReadyForNextTurn; _apiClient = ApiClientFactory.GetClient(); LoadDummyTestCases(); } private void OnDisable() { CommandExecutor.OnQueueUpdated -= Repaint; CommandExecutor.OnPlanUpdated -= HandlePlanUpdated; RequestAnalysisContextCommand.OnAnalysisContextRequested -= HandleAnalysisContextRequested; CommandExecutor.OnContextReadyForNextTurn -= HandleContextReadyForNextTurn; _cancellationTokenSource?.Cancel(); _cancellationTokenSource?.Dispose(); } private void HandlePlanUpdated(List plan) { _currentPlan = plan; Repaint(); } private void HandleAnalysisContextRequested(RequestAnalysisContextParams contextParams) { _isAnalysisContextRequested = true; _requestedRoles = new List(contextParams.subjectRoles); 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(); EditorGUI.BeginDisabledGroup(!SessionManager.HasActiveSession()); DrawContextStagingArea(); EditorGUILayout.Space(); if (!(_apiClient is DummyApiClient)) { DrawSelectionHint(); } EditorGUILayout.Space(); EditorGUILayout.LabelField("Enter your request:"); // --- Conditional UI for Prompt Input --- if (_apiClient is DummyApiClient) { if (!_dummyResponsesLoaded) { EditorGUILayout.HelpBox("Loading dummy responses...", MessageType.Info); } else if (_dummyTriggerPhrases.Any()) { _selectedDummyTriggerIndex = EditorGUILayout.Popup("Select Test Case", _selectedDummyTriggerIndex, _dummyTriggerPhrases.ToArray()); // Update prompt text based on selection to be used by the send function if(_selectedDummyTriggerIndex < _dummyTriggerPhrases.Count) { _promptText = _dummyTriggerPhrases[_selectedDummyTriggerIndex]; } } else { EditorGUILayout.HelpBox("No dummy test cases found in DummyResponses.json.", MessageType.Warning); } } else { // Default text area for other clients like Gemini _promptText = EditorGUILayout.TextArea(_promptText, GUILayout.Height(60)); } // --- End of Conditional UI --- DrawActionButtons(); EditorGUILayout.Space(); DrawPlanArea(); EditorGUILayout.Space(); DrawCommandQueue(); EditorGUI.EndDisabledGroup(); } private void LoadDummyTestCases() { var testCases = DummyApiClient.GetTestCases(); if (testCases != null && testCases.Any()) { _dummyTriggerPhrases = testCases.Select(c => c.TriggerPhrase).ToList(); } else { _dummyTriggerPhrases = new List { "No C# test cases found" }; } _dummyResponsesLoaded = true; Repaint(); } private void DrawSessionManagement() { EditorGUILayout.BeginHorizontal(); if (SessionManager.HasActiveSession()) { EditorGUILayout.LabelField($"Session: {SessionManager.GetCurrentSessionId()}"); if (GUILayout.Button("End Session")) { SessionManager.EndSession(); ClearStagingArea(); _currentPlan.Clear(); } } else { if (GUILayout.Button("Start New Session")) { SessionManager.StartNewSession(); } } EditorGUILayout.EndHorizontal(); } private void DrawContextStagingArea() { EditorGUILayout.LabelField("Manual Context", EditorStyles.boldLabel); EditorGUILayout.HelpBox("You can manually provide context by dragging and dropping any asset (prefab, script, etc.) or GameObject into the slots below.", MessageType.Info); if (GUILayout.Button("Add Manual Context Slot", GUILayout.Width(200))) { _stagedContextObjects.Add(null); if (_isAnalysisContextRequested) _requestedRoles.Add("Custom Role"); } var indexToRemove = -1; for (var i = 0; i < _stagedContextObjects.Count; i++) { EditorGUILayout.BeginHorizontal(); var label = (_isAnalysisContextRequested && i < _requestedRoles.Count) ? _requestedRoles[i] : $"Slot {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) { _stagedContextObjects.RemoveAt(indexToRemove); if (_isAnalysisContextRequested && indexToRemove < _requestedRoles.Count) { _requestedRoles.RemoveAt(indexToRemove); } Repaint(); } } private void DrawSelectionHint() { EditorGUILayout.LabelField("Automatic Context (From Selection)", EditorStyles.boldLabel); var selectedObjects = Selection.objects; if (selectedObjects.Length > 0) { EditorGUILayout.HelpBox("The following selected objects will automatically be included as context.", MessageType.None); EditorGUI.BeginDisabledGroup(true); foreach (var obj in selectedObjects) { EditorGUILayout.ObjectField(obj.name, obj, typeof(Object), true); } EditorGUI.EndDisabledGroup(); } else { EditorGUILayout.HelpBox("No objects are currently selected in the editor.", MessageType.None); } } private void DrawActionButtons() { EditorGUILayout.BeginHorizontal(); EditorGUI.BeginDisabledGroup(_isRequestInProgress); if (_isAnalysisContextRequested) { if (GUILayout.Button("Provide Staged Context to LLM")) { _stagedContextObjects.AddRange(Selection.objects); var contextSummary = ContextBuilder.BuildTier1Summary(_stagedContextObjects); var followUpMessage = "Here is the context you requested:\n" + contextSummary; _ = SendFollowUpAsync(followUpMessage); } } else { if (GUILayout.Button("Send Prompt")) { _ = SendInitialPromptAsync(); } } EditorGUI.EndDisabledGroup(); EditorGUI.BeginDisabledGroup(!_isRequestInProgress); if (GUILayout.Button("Stop", GUILayout.Width(80))) { _cancellationTokenSource?.Cancel(); } EditorGUI.EndDisabledGroup(); EditorGUILayout.EndHorizontal(); } private void DrawPlanArea() { if (_currentPlan == null || !_currentPlan.Any()) return; EditorGUILayout.LabelField("Current Plan", EditorStyles.boldLabel); EditorGUILayout.BeginVertical(EditorStyles.helpBox); for(var i = 0; i < _currentPlan.Count; i++) { var step = _currentPlan[i]; var label = $"{i + 1}. {step}"; EditorGUILayout.LabelField(label, EditorStyles.wordWrappedLabel); } EditorGUILayout.EndVertical(); } private void DrawCommandQueue() { 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); EditorGUI.BeginDisabledGroup(_isRequestInProgress); if (GUILayout.Button("Execute Next Command")) { _ = CommandExecutor.ExecuteNextCommand(); } EditorGUI.EndDisabledGroup(); } else { EditorGUILayout.LabelField("No pending commands."); } EditorGUILayout.EndScrollView(); } private async Task SendInitialPromptAsync() { _isRequestInProgress = true; _cancellationTokenSource = new CancellationTokenSource(); _currentPlan.Clear(); Repaint(); try { var combinedContext = new List(); combinedContext.AddRange(_stagedContextObjects); if (!(_apiClient is DummyApiClient)) { combinedContext.AddRange(Selection.objects); } await _apiClient.SendPrompt(_promptText, combinedContext, _cancellationTokenSource.Token); } catch (TaskCanceledException) { Debug.Log("[MCPWindow] Task was cancelled."); } finally { _isRequestInProgress = false; _cancellationTokenSource.Dispose(); _cancellationTokenSource = null; Repaint(); } } private async Task SendFollowUpAsync(string followUpMessage) { _isRequestInProgress = true; _cancellationTokenSource = new CancellationTokenSource(); Repaint(); try { await _apiClient.SendFollowUp(followUpMessage, _cancellationTokenSource.Token); } catch (TaskCanceledException) { Debug.Log("[MCPWindow] Task was cancelled."); } finally { _isRequestInProgress = false; _cancellationTokenSource.Dispose(); _cancellationTokenSource = null; if (_isAnalysisContextRequested) { ClearStagingArea(); } Repaint(); } } private void ClearStagingArea() { _isAnalysisContextRequested = false; _stagedContextObjects.Clear(); _requestedRoles.Clear(); } } }