terra-llm/Assets/LLM/Editor/Core/CommandExecutor.cs

using System;
using System.Linq;
using System.Threading.Tasks;
using UnityEditor;
using UnityEngine;
using LLM.Editor.Data;
using System.Reflection;
using LLM.Editor.Client;
using LLM.Editor.Commands;
using LLM.Editor.Helper;
using System.Collections.Generic;

namespace LLM.Editor.Core
{
    /// <summary>
    /// Responsible for finding and executing commands from the session queue.
    /// </summary>
    [InitializeOnLoad]
    public static class CommandExecutor
    {
        private static List<CommandData> _commandQueue;
        private static CommandContext _currentContext;
        private static string _currentUserPrompt;
        private static int _initialCommandCount;

        public static Action OnQueueUpdated;
        public static event Action<string> OnContextReadyForNextTurn;
        public static event Action<List<string>> OnPlanUpdated;
        
        private static readonly HashSet<string> AutoExecutableCommands = new()
        {
            nameof(GatherContextCommand),
            nameof(UpdateWorkingContextCommand)
        };

        static CommandExecutor()
        {
            // This will run when the editor loads, including after a recompile.
            EditorApplication.delayCall += Initialize;
        }

        private static void Initialize()
        {
            Debug.Log("[CommandExecutor] Initializing...");
            if (SessionManager.HasActiveSession())
            {
                _commandQueue = SessionManager.LoadCommandQueue();
                if (!_commandQueue.Any())
                {
                    return;
                }
                Debug.Log($"[CommandExecutor] Resuming session with {_commandQueue.Count} command(s) in queue.");
                
                _currentContext = new CommandContext
                {
                    IdentifierMap = SessionManager.LoadIdentifierMap()
                };
                _currentUserPrompt = SessionManager.LoadCurrentUserPrompt();
                
                _initialCommandCount = _commandQueue.Count; 
                OnQueueUpdated?.Invoke();
                TriggerAutoExecution();
            }
            else
            {
                _commandQueue = new List<CommandData>();
            }
        }

        public static void SetQueue(List<CommandData> commands, string userPrompt)
        {
            _commandQueue = commands;
            _currentContext = new CommandContext
            {
                IdentifierMap = SessionManager.LoadIdentifierMap()
            };
            _currentUserPrompt = userPrompt;
            _initialCommandCount = commands.Count; 
            
            SessionManager.SaveCommandQueue(_commandQueue);
            SessionManager.SaveCurrentUserPrompt(_currentUserPrompt);
            
            OnQueueUpdated?.Invoke();
            Debug.Log($"[CommandExecutor] Queue set with {_commandQueue.Count} commands.");
            foreach (var commandData in _commandQueue)
            {
                Debug.Log($"<color=cyan>[CommandExecutor] {commandData.commandName}</color>");
            }
            TriggerAutoExecution();
        }

        private static void ClearQueue()
        {
            _commandQueue.Clear();
            SessionManager.SaveCommandQueue(_commandQueue);
            OnQueueUpdated?.Invoke();
        }

        public static bool HasPendingCommands() => _commandQueue != null &&  _commandQueue.Any();
        
        public static CommandData GetNextCommand() => HasPendingCommands() ? _commandQueue.First() : null;
        
        private static void TriggerAutoExecution()
        {
            EditorApplication.delayCall += async () =>
            {
                if (IsNextCommandAutoExecutable())
                {
                    await ExecuteNextCommand();
                }
            };
        }

        public static async Task ExecuteNextCommand()
        {
            if (!HasPendingCommands())
            {
                Debug.LogWarning("[CommandExecutor] No commands to execute.");
                return;
            }

            var commandData = _commandQueue.First();
            var outcome = CommandOutcome.Error;
            float[] promptEmbedding = null;
            
            try
            {
                EditorApplication.LockReloadAssemblies();
                
                // Generate the embedding *before* executing the command, but only for non-dummy clients.
                var apiClient = ApiClientFactory.GetClient();
                if (apiClient is not DummyApiClient)
                {
                    if (!string.IsNullOrEmpty(_currentUserPrompt))
                    {
                        if (apiClient != null)
                        {
                            promptEmbedding = await EmbeddingHelper.GetEmbedding(_currentUserPrompt, apiClient.GetAuthToken);
                        }
                        else
                        {
                            Debug.LogError("[CommandExecutor] Could not get API client to generate embedding.");
                        }
                    }
                }

                var commandInstance = CreateCommandInstance(commandData);
                if (commandInstance != null)
                {
                    Debug.Log($"[CommandExecutor] Executing: {commandData.commandName}");
                    
                    var planBefore = _currentContext.Plan;
                    outcome = commandInstance.Execute(_currentContext);
                    var planAfter = _currentContext.Plan;

                    if (planAfter != null && !planAfter.SequenceEqual(planBefore ?? Enumerable.Empty<string>()))
                    {
                        OnPlanUpdated?.Invoke(planAfter);
                    }

                    switch (outcome)
                    {
                        case CommandOutcome.Success:
                            _commandQueue.RemoveAt(0);
                            TriggerAutoExecution();
                            break;
                        
                        case CommandOutcome.Error:
                            var errorMessage = $"Command '{commandData.commandName}' failed.";
                            Debug.LogError($"[CommandExecutor] {errorMessage}. Message:\n{_currentContext.ErrorMessage ?? "An unspecified error occurred."}");
                            var errorContext = new 
                            {
                                lastActionStatus = "Error",
                                failedCommand = commandData.commandName,
                                errorMessage = _currentContext.ErrorMessage ?? "An unspecified error occurred."
                            };
                            OnContextReadyForNextTurn?.Invoke(errorContext.ToJson(true));
                            ClearQueue();
                            break;
                        
                        case CommandOutcome.AwaitingNextTurn:
                            Debug.Log("[CommandExecutor] Pausing queue. Awaiting next turn with LLM.");
                            if (_currentContext.CurrentSubject is string detailedContextJson)
                            {
                                OnContextReadyForNextTurn?.Invoke(detailedContextJson);
                            }
                            break;
                    }
                }
                else
                {
                    Debug.LogError($"[CommandExecutor] Could not create instance for command: {commandData.commandName}");
                    ClearQueue();
                }
            }
            catch(Exception e)
            {
                Debug.LogError($"[CommandExecutor] Failed to execute command '{commandData.commandName}'");
                Debug.LogException(e);
                ClearQueue();
            }
            finally
            {
                var record = new InteractionRecord
                {
                    UserPrompt = _currentUserPrompt,
                    LLMResponse = commandData,
                    Outcome = outcome,
                    Feedback = outcome == CommandOutcome.Error ? _currentContext.ErrorMessage : "Success",
                    isMultiStep = _initialCommandCount > 1,
                    PromptEmbedding = promptEmbedding
                };
                MemoryLogger.AddRecord(record);
                
                SessionManager.SaveCommandQueue(_commandQueue);
                OnQueueUpdated?.Invoke();
                EditorApplication.UnlockReloadAssemblies();
            }
        }
        
        private static bool IsNextCommandAutoExecutable()
        {
            if (!HasPendingCommands()) return false;
            var nextCommandName = GetNextCommand().commandName;
            var cleanName = nextCommandName.EndsWith("Command") ? nextCommandName : nextCommandName + "Command";
            return AutoExecutableCommands.Contains(cleanName);
        }

        private static ICommand CreateCommandInstance(CommandData data)
        {   
            var commandClassName = data.commandName.EndsWith("Command") ? data.commandName : $"{data.commandName}Command";
            var type = Assembly.GetExecutingAssembly().GetTypes()
                .FirstOrDefault(t => t.Namespace == "LLM.Editor.Commands" && t.Name == commandClassName);

            if (type != null)
            {
                var jsonString = data.jsonData?.ToString() ?? "{}";
                return (ICommand)Activator.CreateInstance(type, jsonString);
            }

            Debug.LogError($"[CommandExecutor] Command type '{commandClassName}' not found.");
            return null;
        }
    }
}