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

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)
            {
                Debug.LogError("[UpdateWorkingContextCommand] Invalid parameters.");
                return CommandOutcome.Error;
            }

            var workingContext = SessionManager.LoadWorkingContext();

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

                if (_params.updates != null)
                {
                    workingContext.Merge(_params.updates, new JsonMergeSettings
                    {
                        MergeArrayHandling = MergeArrayHandling.Replace
                    });
                }
            }

            SessionManager.SaveWorkingContext(workingContext);
            Debug.Log($"[UpdateWorkingContextCommand] Working context updated:\n{workingContext.ToString(Newtonsoft.Json.Formatting.Indented)}");

            return CommandOutcome.Success;
        }
    }
}