terra-llm/Assets/LLM/Editor/Commands/UpdateWorkingContextCommand.cs

using System.Linq;
using UnityEngine;
using LLM.Editor.Data;
using LLM.Editor.Core;
using LLM.Editor.Helper;
using Newtonsoft.Json.Linq;
using JetBrains.Annotations;

namespace LLM.Editor.Commands
{
    /// <summary>
    /// A command that allows the LLM to write to, update, or clear its own
    /// short-term memory for the current session.
    /// </summary>
    [UsedImplicitly]
    public class UpdateWorkingContextCommand : ICommand
    {
        [System.Serializable]
        private class UpdateWorkingContextParams
        {
            public JObject updates;
            public string[] keysToRemove;
            public bool clearAll;
        }

        private readonly UpdateWorkingContextParams _params;

        public UpdateWorkingContextCommand(string jsonParams)
        {
            _params = jsonParams?.FromJson<UpdateWorkingContextParams>();
        }

        public CommandOutcome Execute(CommandContext context)
        {
            if (_params == null)
            {
                context.ErrorMessage = "Invalid parameters for UpdateWorkingContextCommand.";
                return CommandOutcome.Error;
            }

            var workingContext = SessionManager.LoadWorkingContext();

            if (_params.clearAll)
            {
                workingContext = new JObject();
                context.Plan = null; 
            }
            
            if (_params.keysToRemove != null)
            {
                foreach (var key in _params.keysToRemove)
                {
                    workingContext.Remove(key);
                    if (key == "plan") context.Plan = null;
                }
            }

            if (_params.updates != null)
            {
                if (_params.updates.TryGetValue("plan", out var planToken) && planToken is JArray planArray)
                {
                    context.Plan = planArray.Select(step => step.ToString()).ToList();
                    Debug.Log($"[UpdateWorkingContextCommand] Extracted plan with {context.Plan.Count} steps.");
                }
                
                foreach (var property in _params.updates.Properties())
                {
                    workingContext[property.Name] = property.Value;
                }
            }

            SessionManager.SaveWorkingContext(workingContext);
            Debug.Log($"[UpdateWorkingContextCommand] Working context updated.");

            return CommandOutcome.Success;
        }
    }
}