swarm_memory 2.1.1 → 2.1.3

data/lib/swarm_sdk/node_orchestrator.rb

@@ -18,18 +18,120 @@ module SwarmSDK
   #   )
   #   result = orchestrator.execute("Build auth system")
   class NodeOrchestrator
-    attr_reader :swarm_name, :nodes, :start_node
+    attr_reader :swarm_name, :nodes, :start_node, :agent_definitions, :agent_instance_cache, :scratchpad
+    attr_writer :swarm_id, :config_for_hooks
+    attr_accessor :swarm_registry_config
 
-    def initialize(swarm_name:, agent_definitions:, nodes:, start_node:)
+    def initialize(swarm_name:, agent_definitions:, nodes:, start_node:, swarm_id: nil, scratchpad: :enabled, allow_filesystem_tools: nil)
       @swarm_name = swarm_name
+      @swarm_id = swarm_id
       @agent_definitions = agent_definitions
       @nodes = nodes
       @start_node = start_node
+      @scratchpad = normalize_scratchpad_mode(scratchpad)
+      @allow_filesystem_tools = allow_filesystem_tools
+      @swarm_registry_config = [] # External swarms config (if using composable swarms)
+      @agent_instance_cache = {
+        primary: {}, # { agent_name => Agent::Chat }
+        delegations: {}, # { "delegate@delegator" => Agent::Chat }
+      }
+
+      # Initialize scratchpad storage based on mode
+      case @scratchpad
+      when :enabled
+        # Enabled mode: single scratchpad shared across all nodes
+        @shared_scratchpad_storage = Tools::Stores::ScratchpadStorage.new
+        @node_scratchpads = nil
+      when :per_node
+        # Per-node mode: separate scratchpad per node (lazy initialized)
+        @shared_scratchpad_storage = nil
+        @node_scratchpads = {}
+      when :disabled
+        # Disabled: no storage at all
+        @shared_scratchpad_storage = nil
+        @node_scratchpads = nil
+      end
 
       validate!
       @execution_order = build_execution_order
     end
 
+    # Alias for compatibility with Swarm interface
+    alias_method :name, :swarm_name
+
+    # Get scratchpad storage for a specific node
+    #
+    # Returns the appropriate scratchpad based on mode:
+    # - :enabled - returns the shared scratchpad (same for all nodes)
+    # - :per_node - returns node-specific scratchpad (lazy initialized)
+    # - :disabled - returns nil
+    #
+    # @param node_name [Symbol] Node name
+    # @return [Tools::Stores::ScratchpadStorage, nil] Scratchpad instance or nil if disabled
+    def scratchpad_for(node_name)
+      case @scratchpad
+      when :enabled
+        @shared_scratchpad_storage
+      when :per_node
+        # Lazy initialization per node
+        @node_scratchpads[node_name] ||= Tools::Stores::ScratchpadStorage.new
+      when :disabled
+        nil
+      end
+    end
+
+    # Get all scratchpad storages (for snapshot/restore)
+    #
+    # @return [Hash] { :shared => scratchpad } or { node_name => scratchpad }
+    def all_scratchpads
+      case @scratchpad
+      when :enabled
+        { shared: @shared_scratchpad_storage }
+      when :per_node
+        @node_scratchpads.dup
+      when :disabled
+        {}
+      end
+    end
+
+    # Check if scratchpad is enabled
+    #
+    # @return [Boolean]
+    def scratchpad_enabled?
+      @scratchpad != :disabled
+    end
+
+    # Check if scratchpad is shared between nodes (enabled mode)
+    #
+    # @return [Boolean]
+    def shared_scratchpad?
+      @scratchpad == :enabled
+    end
+
+    # Check if scratchpad is per-node
+    #
+    # @return [Boolean]
+    def per_node_scratchpad?
+      @scratchpad == :per_node
+    end
+
+    # Backward compatibility accessor
+    #
+    # @return [Tools::Stores::ScratchpadStorage, nil]
+    def shared_scratchpad_storage
+      if @scratchpad == :per_node
+        RubyLLM.logger.warn("NodeOrchestrator: Accessing shared_scratchpad_storage in per-node mode. Use scratchpad_for(node_name) instead.")
+      end
+      @shared_scratchpad_storage
+    end
+
+    # Return the lead agent of the start node for CLI compatibility
+    #
+    # @return [Symbol] Lead agent of the start node
+    def lead_agent
+      @nodes[@start_node].lead_agent
+    end
+
     # Execute the node workflow
     #
     # Executes nodes in topological order, passing output from each node
@@ -44,6 +146,9 @@ module SwarmSDK
       results = {}
       @original_prompt = prompt # Store original prompt for NodeContext
 
+      # Set fiber-local execution context for entire workflow
+      Fiber[:execution_id] = generate_execution_id
+
       # Setup logging if block given
       if block_given?
         # Register callback to collect logs and forward to user's block
@@ -56,10 +161,21 @@ module SwarmSDK
         LogStream.emitter = LogCollector
       end
 
-      @execution_order.each do |node_name|
+      # Dynamic execution with support for goto_node
+      execution_index = 0
+      last_result = nil
+
+      while execution_index < @execution_order.size
+        node_name = @execution_order[execution_index]
         node = @nodes[node_name]
         node_start_time = Time.now
 
+        # Set node-specific swarm_id in fiber storage
+        # Mini-swarms will use ||= to inherit execution_id
+        node_swarm_id = @swarm_id ? "#{@swarm_id}/node:#{node_name}" : nil
+        Fiber[:swarm_id] = node_swarm_id
+        Fiber[:parent_swarm_id] = @swarm_id
+
         # Emit node_start event
         emit_node_start(node_name, node)
 
@@ -102,13 +218,26 @@ module SwarmSDK
           # - Exit 2: Halt workflow with error from STDERR (STDOUT ignored)
           transformed = node.transform_input(input_context, current_input: current_input)
 
-          # Check if transformer requested skipping execution
-          # (from Ruby block returning hash OR bash command exit 1)
-          if transformed.is_a?(Hash) && transformed[:skip_execution]
+          # Check for control flow from transformer
+          control_result = handle_transformer_control_flow(
+            transformed: transformed,
+            node_name: node_name,
+            node: node,
+            node_start_time: node_start_time,
+          )
+
+          case control_result[:action]
+          when :halt
+            return control_result[:result]
+          when :goto
+            execution_index = find_node_index(control_result[:target])
+            current_input = control_result[:content]
+            next
+          when :skip
             skip_execution = true
-            skip_content = transformed[:content] || transformed["content"]
-          else
-            current_input = transformed
+            skip_content = control_result[:content]
+          when :continue
+            current_input = control_result[:content]
           end
         end
 
@@ -130,6 +259,9 @@ module SwarmSDK
           mini_swarm = build_swarm_for_node(node)
           result = mini_swarm.execute(current_input)
 
+          # Cache agent instances for context preservation
+          cache_agent_instances(mini_swarm, node)
+
           # If result has error, log it with backtrace
           if result.error
             RubyLLM.logger.error("NodeOrchestrator: Node '#{node_name}' failed: #{result.error.message}")
@@ -138,6 +270,7 @@ module SwarmSDK
         end
 
         results[node_name] = result
+        last_result = result
 
         # Transform output for next node using NodeContext
         output_context = NodeContext.for_output(
@@ -146,7 +279,29 @@ module SwarmSDK
           original_prompt: @original_prompt,
           node_name: node_name,
         )
-        current_input = node.transform_output(output_context)
+        transformed_output = node.transform_output(output_context)
+
+        # Check for control flow from output transformer
+        control_result = handle_output_transformer_control_flow(
+          transformed: transformed_output,
+          node_name: node_name,
+          node: node,
+          node_start_time: node_start_time,
+          skip_execution: skip_execution,
+          result: result,
+        )
+
+        case control_result[:action]
+        when :halt
+          return control_result[:result]
+        when :goto
+          execution_index = find_node_index(control_result[:target])
+          current_input = control_result[:content]
+          emit_node_stop(node_name, node, result, Time.now - node_start_time, skip_execution)
+          next
+        when :continue
+          current_input = control_result[:content]
+        end
 
         # For agent-less nodes, update the result with transformed content
         # This ensures all_results contains the actual output, not the input
@@ -158,22 +313,104 @@ module SwarmSDK
             duration: result.duration,
             error: result.error,
           )
+          last_result = results[node_name]
         end
 
         # Emit node_stop event
         node_duration = Time.now - node_start_time
         emit_node_stop(node_name, node, result, node_duration, skip_execution)
+
+        execution_index += 1
       end
 
-      results.values.last
+      last_result
     ensure
+      # NodeOrchestrator always clears (always sets up logging)
+      Fiber[:execution_id] = nil
+      Fiber[:swarm_id] = nil
+      Fiber[:parent_swarm_id] = nil
+
       # Reset logging state for next execution
       LogCollector.reset!
       LogStream.reset!
     end
 
+    # Create snapshot of current workflow state
+    #
+    # Returns a Snapshot object containing agent conversations, context state,
+    # and scratchpad data from all nodes that have been executed. The snapshot
+    # captures the state of agents in the agent_instance_cache (both primary and
+    # delegation instances), as well as scratchpad storage.
+    #
+    # Configuration (agent definitions, nodes, transformers) stays in your code
+    # and is NOT included in snapshots.
+    #
+    # Scratchpad behavior depends on scratchpad mode:
+    # - :enabled (default): single scratchpad shared across all nodes
+    # - :per_node: separate scratchpad per node
+    # - :disabled: no scratchpad data
+    #
+    # @return [Snapshot] Snapshot object with convenient serialization methods
+    #
+    # @example Save snapshot to JSON file
+    #   orchestrator = NodeOrchestrator.new(...)
+    #   orchestrator.execute("Build feature")
+    #   snapshot = orchestrator.snapshot
+    #   snapshot.write_to_file("workflow_session.json")
+    def snapshot
+      StateSnapshot.new(self).snapshot
+    end
+
+    # Restore workflow state from snapshot
+    #
+    # Accepts a Snapshot object, hash, or JSON string. Validates compatibility
+    # between snapshot and current orchestrator configuration. Restores agent
+    # conversations that exist in the agent_instance_cache.
+    #
+    # The orchestrator must be created with the SAME configuration (agent definitions,
+    # nodes) as when the snapshot was created. Only conversation state is restored.
+    #
+    # For agents with reset_context: false, restored conversations will be injected
+    # during node execution. Agents not in cache yet will be skipped (they haven't
+    # been used yet, so there's nothing to restore).
+    #
+    # @param snapshot [Snapshot, Hash, String] Snapshot object, hash, or JSON string
+    # @return [RestoreResult] Result with warnings about skipped agents
+    #
+    # @example Restore from Snapshot object
+    #   orchestrator = NodeOrchestrator.new(...)  # Same config as snapshot
+    #   snapshot = Snapshot.from_file("workflow_session.json")
+    #   result = orchestrator.restore(snapshot)
+    #   if result.success?
+    #     puts "All agents restored"
+    #   else
+    #     puts result.summary
+    #   end
+    #
+    # Restore orchestrator state from snapshot
+    #
+    # By default, uses current system prompts from agent definitions (YAML + SDK defaults + plugin injections).
+    # Set preserve_system_prompts: true to use historical prompts from snapshot.
+    #
+    # @param snapshot [Snapshot, Hash, String] Snapshot object, hash, or JSON string
+    # @param preserve_system_prompts [Boolean] Use historical system prompts instead of current config (default: false)
+    # @return [RestoreResult] Result with warnings about partial restores
+    def restore(snapshot, preserve_system_prompts: false)
+      StateRestorer.new(self, snapshot, preserve_system_prompts: preserve_system_prompts).restore
+    end
+
     private
 
+    # Generate a unique execution ID for workflow
+    #
+    # Creates an execution ID that uniquely identifies a single orchestrator.execute() call.
+    # Format: "exec_workflow_{random_hex}"
+    #
+    # @return [String] Generated execution ID (e.g., "exec_workflow_a3f2b1c8")
+    def generate_execution_id
+      "exec_workflow_#{SecureRandom.hex(8)}"
+    end
+
     # Emit node_start event
     #
     # @param node_name [Symbol] Name of the node
@@ -284,21 +521,49 @@ module SwarmSDK
     # Creates a new Swarm with only the agents specified in the node,
     # configured with the node's delegation topology.
     #
+    # For agents with reset_context: false, injects cached instances
+    # to preserve conversation history across nodes.
+    #
+    # Scratchpad behavior depends on mode:
+    # - :enabled - all nodes use the same scratchpad instance
+    # - :per_node - each node gets its own scratchpad instance
+    # - :disabled - no scratchpad
+    #
     # @param node [Node::Builder] Node configuration
     # @return [Swarm] Configured swarm instance
     def build_swarm_for_node(node)
-      swarm = Swarm.new(name: "#{@swarm_name}:#{node.name}")
+      # Build hierarchical swarm_id if parent has one (nil auto-generates)
+      node_swarm_id = @swarm_id ? "#{@swarm_id}/node:#{node.name}" : nil
+
+      swarm = Swarm.new(
+        name: "#{@swarm_name}:#{node.name}",
+        swarm_id: node_swarm_id,
+        parent_swarm_id: @swarm_id,
+        scratchpad: scratchpad_for(node.name),
+        scratchpad_mode: :enabled, # Mini-swarms always use enabled (scratchpad instance passed in)
+        allow_filesystem_tools: @allow_filesystem_tools,
+      )
+
+      # Setup swarm registry if external swarms are registered
+      if @swarm_registry_config&.any?
+        registry = SwarmRegistry.new(parent_swarm_id: node_swarm_id || swarm.swarm_id)
+        @swarm_registry_config.each do |reg|
+          registry.register(reg[:name], source: reg[:source], keep_context: reg[:keep_context])
+        end
+        swarm.swarm_registry = registry
+      end
 
       # Add each agent specified in this node
       node.agent_configs.each do |config|
         agent_name = config[:agent]
         delegates_to = config[:delegates_to]
+        tools_override = config[:tools]
 
         # Get global agent definition
         agent_def = @agent_definitions[agent_name]
 
-        # Clone definition with node-specific delegation
-        node_specific_def = clone_with_delegation(agent_def, delegates_to)
+        # Clone definition with node-specific overrides
+        node_specific_def = clone_agent_for_node(agent_def, delegates_to, tools_override)
 
         swarm.add_agent(node_specific_def)
       end
@@ -306,17 +571,26 @@ module SwarmSDK
       # Set lead agent
       swarm.lead = node.lead_agent
 
+      # Inject cached agent instances for context preservation
+      inject_cached_agents(swarm, node)
+
       swarm
     end
 
-    # Clone an agent definition with different delegates_to
+    # Clone an agent definition with node-specific overrides
+    #
+    # Allows overriding delegation and tools per node. This enables:
+    # - Different delegation topology per node
+    # - Different tool sets per workflow stage
     #
     # @param agent_def [Agent::Definition] Original definition
     # @param delegates_to [Array<Symbol>] New delegation targets
-    # @return [Agent::Definition] Cloned definition
-    def clone_with_delegation(agent_def, delegates_to)
+    # @param tools [Array<Symbol>, nil] Tool override (nil = use global agent definition)
+    # @return [Agent::Definition] Cloned definition with overrides
+    def clone_agent_for_node(agent_def, delegates_to, tools)
       config = agent_def.to_h
       config[:delegates_to] = delegates_to
+      config[:tools] = tools if tools # Only override if explicitly set
       Agent::Definition.new(agent_def.name, config)
     end
 
@@ -359,7 +633,8 @@ module SwarmSDK
       if order.size < @nodes.size
         unprocessed = @nodes.keys - order
         raise CircularDependencyError,
-          "Circular dependency detected. Unprocessed nodes: #{unprocessed.join(", ")}"
+          "Circular dependency detected. Unprocessed nodes: #{unprocessed.join(", ")}. " \
+            "Use goto_node in transformers to create loops instead of circular depends_on."
       end
 
       # Verify start_node is in the execution order
@@ -380,5 +655,204 @@ module SwarmSDK
 
       order
     end
+
+    # Handle control flow from input transformer
+    #
+    # @param transformed [String, Hash] Result from transformer
+    # @param node_name [Symbol] Current node name
+    # @param node [Node::Builder] Node configuration
+    # @param node_start_time [Time] Node execution start time
+    # @return [Hash] Control result with :action and relevant data
+    def handle_transformer_control_flow(transformed:, node_name:, node:, node_start_time:)
+      return { action: :continue, content: transformed } unless transformed.is_a?(Hash)
+
+      if transformed[:halt_workflow]
+        # Halt entire workflow
+        halt_result = Result.new(
+          content: transformed[:content],
+          agent: "halted:#{node_name}",
+          logs: [],
+          duration: Time.now - node_start_time,
+        )
+        emit_node_stop(node_name, node, halt_result, Time.now - node_start_time, false)
+        { action: :halt, result: halt_result }
+      elsif transformed[:goto_node]
+        # Jump to different node
+        { action: :goto, target: transformed[:goto_node], content: transformed[:content] }
+      elsif transformed[:skip_execution]
+        # Skip node execution
+        { action: :skip, content: transformed[:content] }
+      else
+        # No control flow - continue normally
+        { action: :continue, content: transformed[:content] }
+      end
+    end
+
+    # Handle control flow from output transformer
+    #
+    # @param transformed [String, Hash] Result from transformer
+    # @param node_name [Symbol] Current node name
+    # @param node [Node::Builder] Node configuration
+    # @param node_start_time [Time] Node execution start time
+    # @param skip_execution [Boolean] Whether node execution was skipped
+    # @param result [Result] Node execution result
+    # @return [Hash] Control result with :action and relevant data
+    def handle_output_transformer_control_flow(transformed:, node_name:, node:, node_start_time:, skip_execution:, result:)
+      # If not a hash, it's just transformed content - continue normally
+      return { action: :continue, content: transformed } unless transformed.is_a?(Hash)
+
+      if transformed[:halt_workflow]
+        # Halt entire workflow
+        halt_result = Result.new(
+          content: transformed[:content],
+          agent: result.agent,
+          logs: result.logs,
+          duration: result.duration,
+        )
+        emit_node_stop(node_name, node, halt_result, Time.now - node_start_time, skip_execution)
+        { action: :halt, result: halt_result }
+      elsif transformed[:goto_node]
+        # Jump to different node
+        { action: :goto, target: transformed[:goto_node], content: transformed[:content] }
+      else
+        # Hash without control flow keys - treat as regular hash with :content key
+        # This handles the case where transformer returns a hash that's not for control flow
+        { action: :continue, content: transformed[:content] || transformed }
+      end
+    end
+
+    # Find the index of a node in the execution order
+    #
+    # @param node_name [Symbol] Node name to find
+    # @return [Integer] Index in execution order
+    # @raise [ConfigurationError] If node not found
+    def find_node_index(node_name)
+      index = @execution_order.index(node_name)
+      unless index
+        raise ConfigurationError,
+          "goto_node target '#{node_name}' not found. Available nodes: #{@execution_order.join(", ")}"
+      end
+      index
+    end
+
+    # Cache agent instances from a swarm for potential reuse
+    #
+    # Only caches agents that have reset_context: false in this node.
+    # This allows preserving conversation history across nodes.
+    #
+    # @param swarm [Swarm] Swarm instance that just executed
+    # @param node [Node::Builder] Node configuration
+    # @return [void]
+    def cache_agent_instances(swarm, node)
+      return unless swarm.agents
+
+      node.agent_configs.each do |config|
+        agent_name = config[:agent]
+        reset_context = config[:reset_context]
+
+        # Only cache if reset_context: false
+        next if reset_context
+
+        # Cache primary agent
+        agent_instance = swarm.agents[agent_name]
+        @agent_instance_cache[:primary][agent_name] = agent_instance if agent_instance
+
+        # V7.0: Cache delegation instances atomically (together with primary)
+        agent_def = @agent_definitions[agent_name]
+        agent_def.delegates_to.each do |delegate_name|
+          delegation_key = "#{delegate_name}@#{agent_name}"
+          delegation_instance = swarm.delegation_instances[delegation_key]
+
+          if delegation_instance
+            @agent_instance_cache[:delegations][delegation_key] = delegation_instance
+          end
+        end
+      end
+    end
+
+    # Inject cached agent instances into a swarm
+    #
+    # For agents with reset_context: false, reuses cached instances to preserve context.
+    # Forces agent initialization first (by accessing .agents), then swaps in cached instances.
+    #
+    # @param swarm [Swarm] Swarm instance to inject into
+    # @param node [Node::Builder] Node configuration
+    # @return [void]
+    def inject_cached_agents(swarm, node)
+      # Check if any agents need context preservation
+      has_preserved = node.agent_configs.any? do |c|
+        !c[:reset_context] && (
+          @agent_instance_cache[:primary][c[:agent]] ||
+          has_cached_delegations_for?(c[:agent])
+        )
+      end
+      return unless has_preserved
+
+      # V7.0 CRITICAL FIX: Force initialization FIRST
+      # Without this, @agents will be replaced by initialize_all, losing our injected instances
+      swarm.agent(node.agent_configs.first[:agent]) # Triggers lazy init
+
+      # Now safely inject cached instances
+      agents_hash = swarm.agents
+      delegation_hash = swarm.delegation_instances
+
+      # Inject cached PRIMARY agents
+      node.agent_configs.each do |config|
+        agent_name = config[:agent]
+        next if config[:reset_context]
+
+        cached_agent = @agent_instance_cache[:primary][agent_name]
+        next unless cached_agent
+
+        # Replace freshly initialized agent with cached instance
+        agents_hash[agent_name] = cached_agent
+      end
+
+      # Inject cached DELEGATION instances (atomic with primary)
+      node.agent_configs.each do |config|
+        agent_name = config[:agent]
+        next if config[:reset_context]
+
+        agent_def = @agent_definitions[agent_name]
+
+        agent_def.delegates_to.each do |delegate_name|
+          delegation_key = "#{delegate_name}@#{agent_name}"
+          cached_delegation = @agent_instance_cache[:delegations][delegation_key]
+          next unless cached_delegation
+
+          # Replace freshly initialized delegation instance
+          # V7.0: Tool references intact - atomic caching preserves object graph
+          delegation_hash[delegation_key] = cached_delegation
+        end
+      end
+    end
+
+    def has_cached_delegations_for?(agent_name)
+      agent_def = @agent_definitions[agent_name]
+      agent_def.delegates_to.any? do |delegate_name|
+        delegation_key = "#{delegate_name}@#{agent_name}"
+        @agent_instance_cache[:delegations][delegation_key]
+      end
+    end
+
+    # Normalize scratchpad mode parameter
+    #
+    # Accepts symbols: :enabled, :per_node, or :disabled
+    #
+    # @param value [Symbol, String] Scratchpad mode (strings from YAML converted to symbols)
+    # @return [Symbol] Normalized mode (:enabled, :per_node, or :disabled)
+    # @raise [ArgumentError] If value is invalid
+    def normalize_scratchpad_mode(value)
+      # Convert strings from YAML to symbols
+      value = value.to_sym if value.is_a?(String)
+
+      case value
+      when :enabled, :per_node, :disabled
+        value
+      else
+        raise ArgumentError,
+          "Invalid scratchpad mode: #{value.inspect}. Use :enabled, :per_node, or :disabled"
+      end
+    end
   end
 end