swarm_memory 2.1.3 → 2.1.4

data/lib/swarm_sdk/state_restorer.rb

@@ -1,15 +1,15 @@
 # frozen_string_literal: true
 
 module SwarmSDK
-  # Restores swarm conversation state from snapshots
+  # Restores swarm/workflow conversation state from snapshots
   #
-  # Unified implementation that works for both Swarm and NodeOrchestrator.
+  # Unified implementation that works for both Swarm and Workflow.
   # Validates compatibility between snapshot and current configuration,
   # restores conversation history, context state, scratchpad contents, and
   # read tracking information.
   #
   # Handles configuration mismatches gracefully by skipping agents that
-  # don't exist in the current swarm and returning warnings in RestoreResult.
+  # don't exist in the current swarm/workflow and returning warnings in RestoreResult.
   #
   # ## System Prompt Handling
   #
@@ -32,12 +32,11 @@ module SwarmSDK
   class StateRestorer
     # Initialize state restorer
     #
-    # @param orchestration [Swarm, NodeOrchestrator] Swarm or orchestrator to restore into
+    # @param orchestration [Swarm, Workflow] Swarm or workflow to restore into
     # @param snapshot [Snapshot, Hash, String] Snapshot object, hash, or JSON string
     # @param preserve_system_prompts [Boolean] If true, use system prompts from snapshot instead of current config (default: false)
     def initialize(orchestration, snapshot, preserve_system_prompts: false)
       @orchestration = orchestration
-      @type = orchestration.is_a?(SwarmSDK::NodeOrchestrator) ? :node_orchestrator : :swarm
       @preserve_system_prompts = preserve_system_prompts
 
       # Handle different input types
@@ -74,7 +73,7 @@ module SwarmSDK
       restore_delegation_conversations(validation.restorable_delegations)
       restore_scratchpad
       restore_read_tracking
-      restore_memory_read_tracking
+      restore_plugin_states
 
       # Phase 3: Return result with warnings
       SwarmSDK::RestoreResult.new(
@@ -91,8 +90,8 @@ module SwarmSDK
     # @raise [StateError] if version is unsupported
     def validate_version!
       version = @snapshot_data[:version] || @snapshot_data["version"]
-      unless version == "1.0.0"
-        raise StateError, "Unsupported snapshot version: #{version}"
+      unless version == "2.1.0"
+        raise StateError, "Unsupported snapshot version: #{version}. Expected: 2.1.0"
       end
     end
 
@@ -100,17 +99,16 @@ module SwarmSDK
     #
     # @raise [StateError] if types don't match
     def validate_type_match!
-      snapshot_type = (@snapshot_data[:type] || @snapshot_data["type"]).to_sym
-      unless snapshot_type == @type
-        raise StateError, "Snapshot type '#{snapshot_type}' doesn't match orchestration type '#{@type}'"
+      snapshot_type = (@snapshot_data[:type] || @snapshot_data["type"]).to_s.downcase
+      actual_type = @orchestration.class.name.split("::").last.downcase
+
+      unless snapshot_type == actual_type
+        raise StateError, "Snapshot type '#{snapshot_type}' doesn't match orchestration type '#{actual_type}'"
       end
     end
 
     # Validate compatibility between snapshot and current configuration
     #
-    # Checks which agents from the snapshot exist in current configuration
-    # and generates warnings for any that don't match.
-    #
     # @return [ValidationResult] Validation results
     def validate_compatibility
       warnings = []
@@ -145,9 +143,7 @@ module SwarmSDK
       delegation_instances&.each do |instance_name, _data|
         base_name, delegator_name = instance_name.split("@")
 
-        # Delegation can be restored if:
-        # 1. The base agent exists in current configuration (may not be in snapshot as primary agent)
-        # 2. The delegator was a restorable primary agent from the snapshot
+        # Delegation can be restored if both agents exist in current configuration
         if current_agents.include?(base_name.to_sym) &&
             restorable_agents.include?(delegator_name.to_sym)
           restorable_delegations << instance_name
@@ -157,7 +153,7 @@ module SwarmSDK
             type: :delegation_instance_not_restorable,
             instance: instance_name,
             message: "Delegation instance '#{instance_name}' cannot be restored " \
-              "(base agent or delegator not in current swarm).",
+              "(base agent or delegator not in current swarm/workflow).",
           }
         end
       end
@@ -173,116 +169,107 @@ module SwarmSDK
 
     # Restore orchestration metadata
     #
-    # For Swarm: restores first_message_sent flag
-    # For NodeOrchestrator: no additional metadata to restore
-    #
     # @return [void]
     def restore_metadata
-      # Restore type-specific metadata
-      if @type == :swarm
-        # Restore first_message_sent flag for Swarm only
-        swarm_data = @snapshot_data[:swarm] || @snapshot_data["swarm"]
-        first_sent = swarm_data[:first_message_sent] || swarm_data["first_message_sent"]
+      # Restore metadata
+      metadata = @snapshot_data[:metadata] || @snapshot_data["metadata"]
+      return unless metadata
+
+      # Restore first_message_sent flag (Swarm only, no-op for Workflow)
+      if @orchestration.respond_to?(:first_message_sent=)
+        first_sent = metadata[:first_message_sent] || metadata["first_message_sent"]
         @orchestration.first_message_sent = first_sent
       end
-      # NodeOrchestrator has no additional metadata to restore
     end
 
     # Restore agent conversations
     #
+    # Uses interface methods - no type checking!
+    #
     # @param restorable_agents [Array<Symbol>] Agents that can be restored
     # @return [void]
     def restore_agent_conversations(restorable_agents)
       restorable_agents.each do |agent_name|
-        # Get agent chat from appropriate source
-        agent_chat = if @type == :swarm
-          # Swarm: agents are lazily initialized, access triggers init
+        # For Swarm: lazy initialization triggers when we call agent()
+        # For Workflow: agents are in cache if already used, otherwise skip
+        agent_chat = if @orchestration.is_a?(Swarm)
           @orchestration.agent(agent_name)
         else
-          # NodeOrchestrator: agents are cached lazily during node execution
-          # If restoring before first execution, cache will be empty
-          # We need to create agents now so they can be injected later
-          cache = @orchestration.agent_instance_cache[:primary]
-          unless cache[agent_name]
-            # For NodeOrchestrator, we can't easily create agents here
-            # because we'd need the full swarm setup (initializer, etc.)
-            # Skip this agent if it's not in cache yet
-            next
-          end
-
-          cache[agent_name]
+          # Workflow: only restore if agent is already in cache
+          @orchestration.primary_agents[agent_name]
         end
 
-        # Get agent snapshot data - handle both symbol and string keys
+        next unless agent_chat
+
+        # Get agent snapshot data
         agents_data = @snapshot_data[:agents] || @snapshot_data["agents"]
         snapshot_data = agents_data[agent_name] || agents_data[agent_name.to_s]
-        next unless snapshot_data # Skip if agent not in snapshot (shouldn't happen due to validation)
-
-        # Clear existing messages FIRST (before adding system prompt)
-        messages = agent_chat.messages
-        messages.clear
-
-        # Determine which system prompt to use
-        # By default, use current prompt from YAML config (allows prompt iteration)
-        # With preserve_system_prompts: true, use historical prompt from snapshot
-        system_prompt = if @preserve_system_prompts
-          # Historical: Use prompt that was active when snapshot was created
-          snapshot_data[:system_prompt] || snapshot_data["system_prompt"]
-        else
-          # Current: Use prompt from current agent definition (default)
-          agent_definition = @orchestration.agent_definitions[agent_name]
-          agent_definition&.system_prompt
-        end
+        next unless snapshot_data
+
+        # Restore conversation
+        restore_agent_conversation(agent_chat, agent_name, snapshot_data)
+      end
+    end
 
-        # Apply system prompt as system message
-        # NOTE: with_instructions adds a system message, so call AFTER clearing
-        agent_chat.with_instructions(system_prompt) if system_prompt
+    # Restore a single agent's conversation
+    #
+    # @param agent_chat [Agent::Chat] Chat instance
+    # @param agent_name [Symbol] Agent name
+    # @param snapshot_data [Hash] Snapshot data for this agent
+    # @return [void]
+    def restore_agent_conversation(agent_chat, agent_name, snapshot_data)
+      # Determine which system prompt to use
+      system_prompt = if @preserve_system_prompts
+        snapshot_data[:system_prompt] || snapshot_data["system_prompt"]
+      else
+        agent_definition = @orchestration.agent_definitions[agent_name]
+        agent_definition&.system_prompt
+      end
 
-        # Restore conversation messages (after system prompt)
-        conversation = snapshot_data[:conversation] || snapshot_data["conversation"]
-        conversation.each do |msg_data|
-          message = deserialize_message(msg_data)
-          messages << message
-        end
+      # Build complete message list including system message
+      all_messages = []
 
-        # Restore context state
-        context_state = snapshot_data[:context_state] || snapshot_data["context_state"]
-        restore_context_state(agent_chat, context_state)
+      # Add system message first if we have a system prompt
+      if system_prompt
+        all_messages << RubyLLM::Message.new(role: :system, content: system_prompt)
       end
+
+      # Add conversation messages
+      conversation = snapshot_data[:conversation] || snapshot_data["conversation"]
+      restored_messages = conversation.map { |msg_data| deserialize_message(msg_data) }
+      all_messages.concat(restored_messages)
+
+      # Replace all messages using proper abstraction
+      agent_chat.replace_messages(all_messages)
+
+      # Restore context state
+      context_state = snapshot_data[:context_state] || snapshot_data["context_state"]
+      restore_context_state(agent_chat, context_state)
     end
 
     # Deserialize a message from snapshot data
     #
-    # Handles Content objects and tool calls properly.
-    #
     # @param msg_data [Hash] Message data from snapshot
     # @return [RubyLLM::Message] Deserialized message
     def deserialize_message(msg_data)
       # Handle Content objects
       content = if msg_data[:content].is_a?(Hash) && (msg_data[:content].key?(:text) || msg_data[:content].key?("text"))
         content_data = msg_data[:content]
-        # Handle both symbol and string keys from JSON
         text = content_data[:text] || content_data["text"]
         attachments = content_data[:attachments] || content_data["attachments"] || []
 
-        # Recreate Content object
-        # NOTE: Attachments are hashes from JSON - RubyLLM::Content constructor handles this
         RubyLLM::Content.new(text, attachments)
       else
-        # Plain string content
         msg_data[:content]
       end
 
-      # Handle tool calls - deserialize from hash array
-      # IMPORTANT: RubyLLM expects tool_calls to be Hash<String, ToolCall>, not Array!
+      # Handle tool calls
       tool_calls_hash = if msg_data[:tool_calls] && !msg_data[:tool_calls].empty?
         msg_data[:tool_calls].each_with_object({}) do |tc_data, hash|
-          # Handle both symbol and string keys from JSON
           id = tc_data[:id] || tc_data["id"]
           name = tc_data[:name] || tc_data["name"]
           arguments = tc_data[:arguments] || tc_data["arguments"] || {}
 
-          # Use ID as hash key (convert to string for consistency)
           hash[id.to_s] = RubyLLM::ToolCall.new(
             id: id,
             name: name,
@@ -308,115 +295,109 @@ module SwarmSDK
     # @param context_state [Hash] Context state data
     # @return [void]
     def restore_context_state(agent_chat, context_state)
-      # Access via public accessors
       context_manager = agent_chat.context_manager
       agent_context = agent_chat.agent_context
 
-      # Restore warning thresholds (Set - add one by one)
+      # Restore warning thresholds
       if context_state[:warning_thresholds_hit] || context_state["warning_thresholds_hit"]
         thresholds_array = context_state[:warning_thresholds_hit] || context_state["warning_thresholds_hit"]
         thresholds_set = agent_context.warning_thresholds_hit
         thresholds_array.each { |t| thresholds_set.add(t) }
       end
 
-      # Restore compression flag using public setter
+      # Restore compression flag
       compression = context_state[:compression_applied] || context_state["compression_applied"]
       context_manager.compression_applied = compression
 
-      # Restore TodoWrite tracking using public setter
+      # Restore TodoWrite tracking
       todowrite_index = context_state[:last_todowrite_message_index] || context_state["last_todowrite_message_index"]
       agent_chat.last_todowrite_message_index = todowrite_index
 
-      # Restore active skill path using public setter
+      # Restore active skill path
       skill_path = context_state[:active_skill_path] || context_state["active_skill_path"]
       agent_chat.active_skill_path = skill_path
     end
 
     # Restore delegation instance conversations
     #
+    # Uses interface methods - no type checking!
+    #
     # @param restorable_delegations [Array<String>] Delegation instances that can be restored
     # @return [void]
     def restore_delegation_conversations(restorable_delegations)
       restorable_delegations.each do |instance_name|
-        # Get delegation chat from appropriate source
-        delegation_chat = if @type == :swarm
-          @orchestration.delegation_instances[instance_name]
-        else
-          cache = @orchestration.agent_instance_cache[:delegations]
-          unless cache[instance_name]
-            # Skip if delegation not in cache yet (NodeOrchestrator)
-            next
-          end
-
-          cache[instance_name]
-        end
+        # Use interface method - works for both!
+        delegation_chat = @orchestration.delegation_instances_hash[instance_name]
         next unless delegation_chat
 
-        # Get delegation snapshot data - handle both symbol and string keys
+        # Get delegation snapshot data
         delegations_data = @snapshot_data[:delegation_instances] || @snapshot_data["delegation_instances"]
         snapshot_data = delegations_data[instance_name.to_sym] || delegations_data[instance_name.to_s] || delegations_data[instance_name]
-        next unless snapshot_data # Skip if delegation not in snapshot (shouldn't happen due to validation)
-
-        # Clear existing messages FIRST (before adding system prompt)
-        messages = delegation_chat.messages
-        messages.clear
+        next unless snapshot_data
 
-        # Determine which system prompt to use
-        # Extract base agent name from delegation instance (e.g., "bob@jarvis" -> "bob")
+        # Extract base agent name
         base_name = instance_name.to_s.split("@").first.to_sym
 
-        # By default, use current prompt from YAML config (allows prompt iteration)
-        # With preserve_system_prompts: true, use historical prompt from snapshot
-        system_prompt = if @preserve_system_prompts
-          # Historical: Use prompt that was active when snapshot was created
-          snapshot_data[:system_prompt] || snapshot_data["system_prompt"]
-        else
-          # Current: Use prompt from current base agent definition (default)
-          agent_definition = @orchestration.agent_definitions[base_name]
-          agent_definition&.system_prompt
-        end
+        # Restore conversation
+        restore_delegation_conversation(delegation_chat, base_name, snapshot_data)
+      end
+    end
 
-        # Apply system prompt as system message
-        # NOTE: with_instructions adds a system message, so call AFTER clearing
-        delegation_chat.with_instructions(system_prompt) if system_prompt
+    # Restore a single delegation's conversation
+    #
+    # @param delegation_chat [Agent::Chat] Chat instance
+    # @param base_name [Symbol] Base agent name
+    # @param snapshot_data [Hash] Snapshot data
+    # @return [void]
+    def restore_delegation_conversation(delegation_chat, base_name, snapshot_data)
+      # Determine which system prompt to use
+      system_prompt = if @preserve_system_prompts
+        snapshot_data[:system_prompt] || snapshot_data["system_prompt"]
+      else
+        agent_definition = @orchestration.agent_definitions[base_name]
+        agent_definition&.system_prompt
+      end
 
-        # Restore conversation messages (after system prompt)
-        conversation = snapshot_data[:conversation] || snapshot_data["conversation"]
-        conversation.each do |msg_data|
-          message = deserialize_message(msg_data)
-          messages << message
-        end
+      # Build complete message list including system message
+      all_messages = []
 
-        # Restore context state
-        context_state = snapshot_data[:context_state] || snapshot_data["context_state"]
-        restore_context_state(delegation_chat, context_state)
+      # Add system message first if we have a system prompt
+      if system_prompt
+        all_messages << RubyLLM::Message.new(role: :system, content: system_prompt)
       end
+
+      # Restore conversation messages
+      conversation = snapshot_data[:conversation] || snapshot_data["conversation"]
+      restored_messages = conversation.map { |msg_data| deserialize_message(msg_data) }
+      all_messages.concat(restored_messages)
+
+      # Replace all messages using proper abstraction
+      delegation_chat.replace_messages(all_messages)
+
+      # Restore context state
+      context_state = snapshot_data[:context_state] || snapshot_data["context_state"]
+      restore_context_state(delegation_chat, context_state)
     end
 
     # Restore scratchpad contents
     #
-    # For Swarm: uses scratchpad_storage (flat format)
-    # For NodeOrchestrator: { shared: bool, data: ... }
-    #   - shared: true  → :enabled mode (shared across nodes)
-    #   - shared: false → :per_node mode (isolated per node)
-    #
     # @return [void]
     def restore_scratchpad
       scratchpad_data = @snapshot_data[:scratchpad] || @snapshot_data["scratchpad"]
       return unless scratchpad_data&.any?
 
-      if @type == :node_orchestrator
-        restore_node_orchestrator_scratchpad(scratchpad_data)
+      if @orchestration.is_a?(Workflow)
+        restore_workflow_scratchpad(scratchpad_data)
       else
         restore_swarm_scratchpad(scratchpad_data)
       end
     end
 
-    # Restore scratchpad for NodeOrchestrator
+    # Restore scratchpad for Workflow
     #
     # @param scratchpad_data [Hash] { shared: bool, data: ... }
     # @return [void]
-    def restore_node_orchestrator_scratchpad(scratchpad_data)
+    def restore_workflow_scratchpad(scratchpad_data)
       snapshot_shared_mode = scratchpad_data[:shared] || scratchpad_data["shared"]
       data = scratchpad_data[:data] || scratchpad_data["data"]
 
@@ -464,27 +445,31 @@ module SwarmSDK
       read_tracking_data = @snapshot_data[:read_tracking] || @snapshot_data["read_tracking"]
       return unless read_tracking_data
 
-      # Restore tracking for each agent using new API
-      # read_tracking_data format: { agent_name => { file_path => digest } }
       read_tracking_data.each do |agent_name, files_with_digests|
         agent_sym = agent_name.to_sym
         Tools::Stores::ReadTracker.restore_read_files(agent_sym, files_with_digests)
       end
     end
 
-    # Restore memory read tracking state
+    # Restore plugin-specific state for all plugins
     #
     # @return [void]
-    def restore_memory_read_tracking
-      memory_tracking_data = @snapshot_data[:memory_read_tracking] || @snapshot_data["memory_read_tracking"]
-      return unless memory_tracking_data
-      return unless defined?(SwarmMemory::Core::StorageReadTracker)
-
-      # Restore tracking for each agent using new API
-      # memory_tracking_data format: { agent_name => { entry_path => digest } }
-      memory_tracking_data.each do |agent_name, entries_with_digests|
-        agent_sym = agent_name.to_sym
-        SwarmMemory::Core::StorageReadTracker.restore_read_entries(agent_sym, entries_with_digests)
+    def restore_plugin_states
+      plugin_states_data = @snapshot_data[:plugin_states] || @snapshot_data["plugin_states"]
+      return unless plugin_states_data
+
+      plugin_states_data.each do |plugin_name, agents_state|
+        # Find plugin by name
+        plugin = PluginRegistry.all.find { |p| p.name.to_s == plugin_name.to_s }
+        next unless plugin
+
+        # Restore state for each agent
+        agents_state.each do |agent_name, state|
+          agent_sym = agent_name.to_sym
+          # Symbolize keys for consistent access
+          symbolized_state = state.is_a?(Hash) ? state.transform_keys(&:to_sym) : state
+          plugin.restore_agent_state(agent_sym, symbolized_state)
+        end
       end
     end
   end