claude_swarm 1.0.4 → 1.0.6

data/lib/swarm_sdk/node_orchestrator.rb

@@ -20,16 +20,28 @@ module SwarmSDK
   class NodeOrchestrator
     attr_reader :swarm_name, :nodes, :start_node
 
-    def initialize(swarm_name:, agent_definitions:, nodes:, start_node:)
+    def initialize(swarm_name:, agent_definitions:, nodes:, start_node:, scratchpad_enabled: true)
       @swarm_name = swarm_name
       @agent_definitions = agent_definitions
       @nodes = nodes
       @start_node = start_node
+      @scratchpad_enabled = scratchpad_enabled
+      @agent_instance_cache = {} # Cache for preserving agent context across nodes
 
       validate!
       @execution_order = build_execution_order
     end
 
+    # Alias for compatibility with Swarm interface
+    alias_method :name, :swarm_name
+
+    # 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
@@ -56,7 +68,12 @@ 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
 
@@ -102,13 +119,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 +160,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 +171,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 +180,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,14 +214,17 @@ 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
       # Reset logging state for next execution
       LogCollector.reset!
@@ -284,10 +343,18 @@ 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.
+    #
+    # Inherits scratchpad_enabled setting from NodeOrchestrator.
+    #
     # @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}")
+      swarm = Swarm.new(
+        name: "#{@swarm_name}:#{node.name}",
+        scratchpad_enabled: @scratchpad_enabled,
+      )
 
       # Add each agent specified in this node
       node.agent_configs.each do |config|
@@ -306,6 +373,9 @@ module SwarmSDK
       # Set lead agent
       swarm.lead = node.lead_agent
 
+      # Inject cached agent instances for context preservation
+      inject_cached_agents(swarm, node)
+
       swarm
     end
 
@@ -359,7 +429,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 +451,141 @@ 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 # Only cache if agents were initialized
+
+      node.agent_configs.each do |config|
+        agent_name = config[:agent]
+        reset_context = config[:reset_context]
+
+        # Only cache if reset_context is false
+        next if reset_context
+
+        # Cache the agent instance
+        agent_instance = swarm.agents[agent_name]
+        @agent_instance_cache[agent_name] = agent_instance if agent_instance
+      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_agents = node.agent_configs.any? { |c| !c[:reset_context] && @agent_instance_cache[c[:agent]] }
+      return unless has_preserved_agents
+
+      # Force agent initialization by accessing .agents (triggers lazy init)
+      # Then inject cached instances
+      agents_hash = swarm.agents
+
+      node.agent_configs.each do |config|
+        agent_name = config[:agent]
+        reset_context = config[:reset_context]
+
+        # Skip if reset_context is true (want fresh instance)
+        next if reset_context
+
+        # Check if we have a cached instance
+        cached_agent = @agent_instance_cache[agent_name]
+        next unless cached_agent
+
+        # Inject the cached instance (replace the freshly initialized one)
+        agents_hash[agent_name] = cached_agent
+      end
+    end
   end
 end

data/lib/swarm_sdk/plugin.rb

@@ -7,7 +7,32 @@ module SwarmSDK
   # Plugins are self-registering - they call SwarmSDK::PluginRegistry.register
   # when the gem is loaded.
   #
-  # @example Implementing a plugin
+  # ## Adding Custom Attributes to Agents
+  #
+  # Plugins can add custom attributes to Agent::Definition that are preserved
+  # when agents are cloned (e.g., in NodeOrchestrator). To do this:
+  #
+  # 1. Add attr_reader to Agent::Definition for your attribute
+  # 2. Parse the attribute in Agent::Definition#initialize
+  # 3. Implement serialize_config to preserve it during serialization
+  #
+  # @example Plugin with custom agent attributes
+  #   # 1. Extend Agent::Definition (in your plugin gem)
+  #   module SwarmSDK
+  #     module Agent
+  #       class Definition
+  #         attr_reader :my_custom_config
+  #
+  #         alias_method :original_initialize, :initialize
+  #         def initialize(name, config = {})
+  #           @my_custom_config = config[:my_custom_config]
+  #           original_initialize(name, config)
+  #         end
+  #       end
+  #     end
+  #   end
+  #
+  #   # 2. Implement plugin with serialize_config
   #   class MyPlugin < SwarmSDK::Plugin
   #     def name
   #       :my_plugin
@@ -20,9 +45,34 @@ module SwarmSDK
   #     def create_tool(tool_name, context)
   #       # Create and return tool instance
   #     end
+  #
+  #     # Preserve custom config when agents are cloned
+  #     def serialize_config(agent_definition:)
+  #       return {} unless agent_definition.my_custom_config
+  #
+  #       { my_custom_config: agent_definition.my_custom_config }
+  #     end
   #   end
   #
   #   SwarmSDK::PluginRegistry.register(MyPlugin.new)
+  #
+  # Now agents can use your custom config:
+  #
+  #   agent :researcher do
+  #     my_custom_config { option: "value" }
+  #   end
+  #
+  # And it will be preserved when NodeOrchestrator clones the agent!
+  #
+  # @example Real-world: SwarmMemory plugin
+  #   # SwarmMemory adds 'memory' attribute to agents
+  #   class SDKPlugin < SwarmSDK::Plugin
+  #     def serialize_config(agent_definition:)
+  #       return {} unless agent_definition.memory
+  #       { memory: agent_definition.memory }
+  #     end
+  #   end
+  #
   class Plugin
     # Plugin name (must be unique)
     #
@@ -143,5 +193,27 @@ module SwarmSDK
     def on_user_message(agent_name:, prompt:, is_first_message:)
       []
     end
+
+    # Contribute to agent serialization (optional)
+    #
+    # Called when Agent::Definition.to_h is invoked (e.g., for cloning agents
+    # in NodeOrchestrator). Plugins can return config keys that should be
+    # included in the serialized hash to preserve their state.
+    #
+    # This allows plugins to maintain their configuration when agents are
+    # cloned or serialized, without SwarmSDK needing to know about plugin-specific fields.
+    #
+    # @param agent_definition [Agent::Definition] Agent definition
+    # @return [Hash] Config keys to include in to_h (e.g., { memory: config })
+    #
+    # @example Memory plugin serialization
+    #   def serialize_config(agent_definition:)
+    #     return {} unless agent_definition.memory
+    #
+    #     { memory: agent_definition.memory }
+    #   end
+    def serialize_config(agent_definition:)
+      {}
+    end
   end
 end

data/lib/swarm_sdk/prompts/base_system_prompt.md.erb

@@ -69,139 +69,15 @@ When making changes to files, first understand the file's conventions. Mimic exi
 - When you edit something, first look at the surrounding context (especially imports/requires) to understand the choice of frameworks and libraries. Then consider how to make the given change in a way that is most consistent with existing patterns.
 - Always follow security best practices. Never introduce code that exposes or logs secrets and keys. Never commit secrets or keys to repositories.
 
-# Task Management
-
-You have access to the TodoWrite tool to help you manage and plan tasks. Use this tool VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress.
-This tool is also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable.
-
-**CRITICAL WORKFLOW**: When starting a multi-step task:
-1. **FIRST**: Analyze what needs to be done (search, read files, understand scope)
-2. **SECOND**: Create a COMPLETE todo list with ALL known tasks before starting work
-3. **THIRD**: Begin executing tasks, marking them in_progress → completed as you work
-4. **ONLY add new todos** if you discover unexpected work during implementation
-
-**CRITICAL RULES FOR TODO COMPLETION**:
-- Mark EACH task as completed IMMEDIATELY after finishing it (do not batch updates)
-- You MUST complete ALL pending todos before giving your final answer to the user
-- If a task becomes irrelevant, remove it from the list or mark it completed with a note
-- NEVER leave in_progress or pending tasks when you finish responding to the user
-- Before giving your final response, verify all todos are marked completed
-
-Examples:
-
-<example>
-user: Run the build and fix any type errors
-assistant: I'll run the build first to identify all type errors, then create a complete todo list.
-
-[Runs build and finds 3 type errors in 3 different files]
-
-Now I'll create a complete todo list with all the work:
-
-[Uses TodoWrite to create full list:]
-1. Fix type error in auth.ts:45 (in_progress)
-2. Fix type error in user.ts:23 (pending)
-3. Fix type error in api.ts:67 (pending)
-4. Run build again to verify all fixes (pending)
-
-Starting with the first error in auth.ts...
-[Fixes auth.ts error]
-
-[Updates TodoWrite - marks task 1 completed, task 2 in_progress]
-
-Now fixing user.ts...
-[Fixes user.ts error]
-
-[Updates TodoWrite - marks task 2 completed, task 3 in_progress]
-
-Now fixing api.ts...
-[Fixes api.ts error]
-
-[Updates TodoWrite - marks task 3 completed, task 4 in_progress]
-
-Running build to verify...
-[Runs build - all errors fixed!]
-
-[Updates TodoWrite - marks task 4 completed]
-
-All type errors have been fixed and the build passes successfully.
-</example>
-In this example, the assistant: (1) analyzed scope first, (2) created COMPLETE todo list, (3) marked EACH task completed immediately after finishing it, (4) completed ALL tasks before giving final response.
-
-<example>
-user: Help me write a new feature that allows users to track their usage metrics and export them to various formats
-
-A: I'll help you implement a usage metrics tracking and export feature. Let me first research the codebase to understand the scope, then create a complete todo list.
-
-[Searches for existing metrics/telemetry code, examines architecture, identifies integration points]
-
-Based on my analysis, here's the complete implementation plan:
-
-[Uses TodoWrite to create full list:]
-1. Create MetricsCollector class in src/metrics/ (in_progress)
-2. Add database schema for metrics storage (pending)
-3. Implement CSV export format (pending)
-4. Implement JSON export format (pending)
-5. Create MetricsService API endpoints (pending)
-6. Add frontend UI for export (pending)
-7. Write unit tests for collectors and exporters (pending)
-8. Update documentation (pending)
-
-Now starting with the MetricsCollector class...
-[Implements MetricsCollector]
-
-[Updates TodoWrite - marks task 1 completed, task 2 in_progress]
-
-Adding database schema...
-[Creates schema]
-
-[Updates TodoWrite - marks task 2 completed, task 3 in_progress]
-
-Implementing CSV export...
-[Implements CSV exporter]
-
-[Updates TodoWrite - marks task 3 completed, task 4 in_progress]
-
-Implementing JSON export...
-[Implements JSON exporter]
-
-[Updates TodoWrite - marks task 4 completed, task 5 in_progress]
-
-Creating MetricsService API endpoints...
-[Implements API endpoints]
-
-[Updates TodoWrite - marks task 5 completed, task 6 in_progress]
-
-Adding frontend UI for export...
-[Implements UI components]
-
-[Updates TodoWrite - marks task 6 completed, task 7 in_progress]
-
-Writing unit tests...
-[Writes and runs tests - all pass]
-
-[Updates TodoWrite - marks task 7 completed, task 8 in_progress]
-
-Updating documentation...
-[Updates docs with usage examples]
-
-[Updates TodoWrite - marks task 8 completed]
-
-The metrics tracking and export feature is now complete. Users can collect metrics and export them to CSV or JSON formats through both the API and the frontend UI.
-</example>
-
 # Doing tasks
 
 The user will primarily request you perform tasks. This includes solving problems, adding new functionality, refactoring, explaining content, and more. For these tasks the following steps are recommended:
 
-- Use the TodoWrite tool to plan the task if required
 - Use the available search tools to understand the context and the user's query. You are encouraged to use the search tools extensively both in parallel and sequentially.
 - Implement the solution using all tools available to you
-- Mark each todo completed IMMEDIATELY after finishing it
 - Verify the solution if possible with tests. NEVER assume specific test framework or test script. Check the project documentation or search to determine the testing approach.
 - When you have completed a task, if there are linting or validation commands available to you, run them to ensure your work is correct. NEVER assume what these commands are - check the project documentation first.
 NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise the user will feel that you are being too proactive.
-- Before giving your final response: Ensure ALL todos are marked completed. NEVER leave pending or in_progress tasks.
-- IMPORTANT: Always use the TodoWrite tool to plan and track tasks throughout the conversation.
 
 # Tool usage policy
 
@@ -211,8 +87,6 @@ NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTAN
 - If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to delegate a task to multiple agents in parallel, send a single message with multiple DelegateTask tool calls.
 - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit/MultiEdit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead.
 
-IMPORTANT: Always use the TodoWrite tool to plan and track tasks throughout the conversation.
-
 You MUST answer concisely with fewer than 4 lines of text (not including tool use or code generation), unless user asks for detail.
 
 

data/lib/swarm_sdk/result.rb

@@ -2,17 +2,17 @@
 
 module SwarmSDK
   class Result
-    attr_reader :content, :agent, :cost, :tokens, :duration, :logs, :error, :metadata
+    attr_reader :content, :agent, :duration, :logs, :error, :metadata
 
-    def initialize(content: nil, agent:, cost: 0.0, tokens: {}, duration: 0.0, logs: [], error: nil, metadata: {})
+    def initialize(content: nil, agent:, cost: nil, tokens: nil, duration: 0.0, logs: [], error: nil, metadata: {})
       @content = content
       @agent = agent
-      @cost = cost
-      @tokens = tokens
       @duration = duration
       @logs = logs
       @error = error
       @metadata = metadata
+      # Legacy parameters kept for backward compatibility but not stored
+      # Use total_cost and tokens methods instead which calculate from logs
     end
 
     def success?
@@ -23,12 +23,38 @@ module SwarmSDK
       !success?
     end
 
+    # Calculate total cost from logs
+    #
+    # Delegates to total_cost for consistency. This attribute is calculated
+    # dynamically rather than stored.
+    #
+    # @return [Float] Total cost in dollars
+    def cost
+      total_cost
+    end
+
+    # Get token breakdown from logs
+    #
+    # Returns input and output tokens from the last log entry with usage data.
+    # This attribute is calculated dynamically rather than stored.
+    #
+    # @return [Hash] Token breakdown with :input and :output keys, or empty hash if no usage data
+    def tokens
+      last_entry = @logs.reverse.find { |entry| entry.dig(:usage, :cumulative_input_tokens) }
+      return {} unless last_entry
+
+      {
+        input: last_entry.dig(:usage, :cumulative_input_tokens) || 0,
+        output: last_entry.dig(:usage, :cumulative_output_tokens) || 0,
+      }
+    end
+
     def to_h
       {
         content: @content,
         agent: @agent,
-        cost: @cost,
-        tokens: @tokens,
+        cost: cost,
+        tokens: tokens,
         duration: @duration,
         success: success?,
         error: @error&.message,

data/lib/swarm_sdk/swarm/builder.rb

@@ -380,6 +380,7 @@ module SwarmSDK
           agent_definitions: agent_definitions,
           nodes: @nodes,
           start_node: @start_node,
+          scratchpad_enabled: @scratchpad_enabled,
         )
       end