swarm_memory 2.1.2 → 2.1.4

data/lib/swarm_sdk/agent/{chat → chat_helpers}/context_tracker.rb

@@ -2,7 +2,7 @@
 
 module SwarmSDK
   module Agent
-    class Chat < RubyLLM::Chat
+    module ChatHelpers
       # Manages context tracking, delegation tracking, and logging callbacks
       #
       # Responsibilities:
@@ -13,6 +13,25 @@ module SwarmSDK
       # - Check context warnings
       #
       # This is a stateful helper that's instantiated per Agent::Chat instance.
+      #
+      # ## Thread Safety and Fiber-Local Storage
+      #
+      # IMPORTANT: LogStream.emit calls in this class DO NOT explicitly pass
+      # swarm_id, parent_swarm_id, or execution_id. These values are automatically
+      # injected from Fiber-local storage (Fiber[:swarm_id], etc.) by LogStream.emit.
+      #
+      # Why: In threaded environments (Puma, Sidekiq), swarm/agent instances may be
+      # reused across multiple requests/jobs. If we explicitly pass @agent_context.swarm_id,
+      # callbacks would use STALE values from the first request, causing events to be
+      # lost or misattributed.
+      #
+      # By relying on Fiber-local storage, each request/job gets the correct context
+      # even when reusing the same swarm instance. Fiber storage is set at the start
+      # of Swarm#execute and inherited by child fibers (tool calls, delegations).
+      #
+      # This design works correctly in both:
+      # - Single-threaded environments (rails runner, console)
+      # - Multi-threaded environments (Puma, Sidekiq)
       class ContextTracker
         include LoggingHelpers
 
@@ -44,56 +63,19 @@ module SwarmSDK
 
         # Extract agent name from delegation tool name
         #
-        # Converts "DelegateTaskTo[AgentName]" to "agent_name"
-        # Example: "DelegateTaskToWorker" -> "worker"
+        # Converts "#{Tools::Delegate::TOOL_NAME_PREFIX}[AgentName]" to "agent_name"
+        # Example: "WorkWithWorker" -> "worker"
         #
         # @param tool_name [String] Delegation tool name
         # @return [String] Agent name
         def extract_delegate_agent_name(tool_name)
-          # Remove "DelegateTaskTo" prefix and lowercase first letter
-          agent_name = tool_name.to_s.sub(/^DelegateTaskTo/, "")
+          # Remove tool name prefix and lowercase first letter
+          agent_name = tool_name.to_s.sub(/^#{Tools::Delegate::TOOL_NAME_PREFIX}/, "")
           # Convert from PascalCase to lowercase (e.g., "Worker" -> "worker", "BackendDev" -> "backendDev")
           agent_name[0] = agent_name[0].downcase unless agent_name.empty?
           agent_name
         end
 
-        # Check if context usage has crossed warning thresholds and emit warnings
-        #
-        # This should be called after each LLM response to check if we've crossed
-        # any warning thresholds (80%, 90%, etc.)
-        #
-        # @return [void]
-        def check_context_warnings
-          current_percentage = @chat.context_usage_percentage
-
-          Context::CONTEXT_WARNING_THRESHOLDS.each do |threshold|
-            # Only warn once per threshold
-            next if @agent_context.warning_threshold_hit?(threshold)
-            next if current_percentage < threshold
-
-            # Mark threshold as hit and emit warning
-            @agent_context.hit_warning_threshold?(threshold)
-
-            # Trigger automatic compression at 60% threshold
-            if threshold == Context::COMPRESSION_THRESHOLD
-              trigger_automatic_compression
-            end
-
-            LogStream.emit(
-              type: "context_limit_warning",
-              agent: @agent_context.name,
-              model: @chat.model.id,
-              threshold: "#{threshold}%",
-              current_usage: "#{current_percentage}%",
-              tokens_used: @chat.cumulative_total_tokens,
-              tokens_remaining: @chat.tokens_remaining,
-              context_limit: @chat.context_limit,
-              metadata: @agent_context.metadata,
-              compression_triggered: threshold == Context::COMPRESSION_THRESHOLD,
-            )
-          end
-        end
-
         private
 
         # Extract usage information from an assistant message
@@ -107,6 +89,9 @@ module SwarmSDK
               cumulative_input_tokens: @chat.cumulative_input_tokens,
               cumulative_output_tokens: @chat.cumulative_output_tokens,
               cumulative_total_tokens: @chat.cumulative_total_tokens,
+              cumulative_cached_tokens: @chat.cumulative_cached_tokens,
+              cumulative_cache_creation_tokens: @chat.cumulative_cache_creation_tokens,
+              effective_input_tokens: @chat.effective_input_tokens,
               context_limit: @chat.context_limit,
               tokens_used_percentage: "#{@chat.context_usage_percentage}%",
               tokens_remaining: @chat.tokens_remaining,
@@ -118,6 +103,8 @@ module SwarmSDK
           {
             input_tokens: message.input_tokens,
             output_tokens: message.output_tokens,
+            cached_tokens: message.cached_tokens,
+            cache_creation_tokens: message.cache_creation_tokens,
             total_tokens: (message.input_tokens || 0) + (message.output_tokens || 0),
             input_cost: cost_info[:input_cost],
             output_cost: cost_info[:output_cost],
@@ -154,6 +141,10 @@ module SwarmSDK
                 # Final response (finish_reason: "stop") - fire agent_stop
                 trigger_agent_stop(message, tool_executions: @tool_executions)
               end
+
+              # Check context warnings after each assistant message
+              # Uses unified implementation in HookIntegration
+              @chat.check_context_warnings if @chat.respond_to?(:check_context_warnings)
             when :tool
               # Handle delegation tracking and logging (technical plumbing)
               if @agent_context.delegation?(call_id: message.tool_call_id)
@@ -272,43 +263,6 @@ module SwarmSDK
           )
         end
       end
-
-      # Trigger automatic message compression
-      #
-      # Called when context usage crosses 60% threshold. Compresses old tool
-      # results to save context window space while preserving accuracy.
-      #
-      # @return [void]
-      def trigger_automatic_compression
-        return unless @chat.respond_to?(:context_manager)
-
-        # Calculate tokens before compression
-        tokens_before = @chat.cumulative_total_tokens
-
-        # Get compressed messages from ContextManager
-        compressed = @chat.context_manager.auto_compress_on_threshold(@chat.messages, keep_recent: 10)
-
-        # Count how many messages were actually compressed
-        messages_compressed = compressed.count do |msg|
-          msg.content.to_s.include?("[truncated for context management]")
-        end
-
-        # Replace messages array with compressed version
-        @chat.messages.clear
-        compressed.each { |msg| @chat.messages << msg }
-
-        # Log compression event
-        LogStream.emit(
-          type: "context_compression",
-          agent: @agent_context.name,
-          total_messages: @chat.messages.size,
-          messages_compressed: messages_compressed,
-          tokens_before: tokens_before,
-          current_usage: "#{@chat.context_usage_percentage}%",
-          compression_strategy: "progressive_tool_result_compression",
-          keep_recent: 10,
-        ) if LogStream.enabled?
-      end
     end
   end
 end

data/lib/swarm_sdk/agent/chat_helpers/event_emitter.rb

@@ -0,0 +1,204 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+module SwarmSDK
+  module Agent
+    module ChatHelpers
+      # Minimal event emitter that mirrors RubyLLM::Chat's callback pattern
+      #
+      # Provides multi-subscriber support for events like tool_call, tool_result,
+      # new_message, end_message. This is thread-safe and supports unsubscription.
+      module EventEmitter
+        # Represents an active subscription to a callback event.
+        # Returned by {#subscribe} and can be used to unsubscribe later.
+        class Subscription
+          attr_reader :tag
+
+          def initialize(callback_list, callback, monitor:, tag: nil)
+            @callback_list = callback_list
+            @callback = callback
+            @monitor = monitor
+            @tag = tag
+            @active = true
+          end
+
+          # Removes this subscription from the callback list.
+          # @return [Boolean] true if successfully unsubscribed, false if already inactive
+          def unsubscribe # rubocop:disable Naming/PredicateMethod
+            @monitor.synchronize do
+              return false unless @active
+
+              @callback_list.delete(@callback)
+              @active = false
+            end
+            true
+          end
+
+          # Checks if this subscription is still active.
+          # @return [Boolean] true if still subscribed
+          def active?
+            @monitor.synchronize do
+              @active && @callback_list.include?(@callback)
+            end
+          end
+
+          def inspect
+            "#<#{self.class.name} tag=#{@tag.inspect} active=#{active?}>"
+          end
+        end
+
+        # Initialize the event emitter system
+        #
+        # Sets up @callbacks hash and @callback_monitor for thread safety.
+        # Must be called in Chat#initialize.
+        #
+        # @return [void]
+        def initialize_event_emitter
+          @callbacks = {
+            new_message: [],
+            end_message: [],
+            tool_call: [],
+            tool_result: [],
+          }
+          @callback_monitor = Monitor.new
+        end
+
+        # Subscribes to an event with the given block.
+        # Returns a {Subscription} that can be used to unsubscribe.
+        #
+        # @param event [Symbol] The event to subscribe to
+        # @param tag [String, nil] Optional tag for debugging/identification
+        # @yield The block to call when the event fires
+        # @return [Subscription] An object that can be used to unsubscribe
+        # @raise [ArgumentError] if event is not recognized
+        def subscribe(event, tag: nil, &block)
+          @callback_monitor.synchronize do
+            unless @callbacks.key?(event)
+              raise ArgumentError, "Unknown event: #{event}. Valid events: #{@callbacks.keys.join(", ")}"
+            end
+
+            @callbacks[event] << block
+            Subscription.new(@callbacks[event], block, monitor: @callback_monitor, tag: tag)
+          end
+        end
+
+        # Subscribes to an event that automatically unsubscribes after firing once.
+        #
+        # @param event [Symbol] The event to subscribe to
+        # @param tag [String, nil] Optional tag for debugging/identification
+        # @yield The block to call when the event fires (once)
+        # @return [Subscription] An object that can be used to unsubscribe before it fires
+        def once(event, tag: nil, &block)
+          subscription = nil
+          wrapper = lambda do |*args|
+            subscription&.unsubscribe
+            block.call(*args)
+          end
+          subscription = subscribe(event, tag: tag, &wrapper)
+        end
+
+        # Registers a callback for when a new message starts being generated.
+        # Multiple callbacks can be registered and all will fire in registration order.
+        #
+        # @yield Block called when a new message starts
+        # @return [self] for chaining
+        def on_new_message(&block)
+          subscribe(:new_message, &block)
+          self
+        end
+
+        # Registers a callback for when a message is complete.
+        # Multiple callbacks can be registered and all will fire in registration order.
+        #
+        # @yield [Message] Block called with the completed message
+        # @return [self] for chaining
+        def on_end_message(&block)
+          subscribe(:end_message, &block)
+          self
+        end
+
+        # Registers a callback for when a tool is called.
+        # Multiple callbacks can be registered and all will fire in registration order.
+        #
+        # @yield [ToolCall] Block called with the tool call object
+        # @return [self] for chaining
+        def on_tool_call(&block)
+          subscribe(:tool_call, &block)
+          self
+        end
+
+        # Registers a callback for when a tool returns a result.
+        # Multiple callbacks can be registered and all will fire in registration order.
+        #
+        # @yield [Object] Block called with the tool result
+        # @return [self] for chaining
+        def on_tool_result(&block)
+          subscribe(:tool_result, &block)
+          self
+        end
+
+        # Clears all callbacks for the specified event, or all events if none specified.
+        #
+        # @param event [Symbol, nil] The event to clear callbacks for, or nil for all events
+        # @return [self] for chaining
+        def clear_callbacks(event = nil)
+          @callback_monitor.synchronize do
+            if event
+              @callbacks[event]&.clear
+            else
+              @callbacks.each_value(&:clear)
+            end
+          end
+          self
+        end
+
+        # Returns the number of callbacks registered for the specified event.
+        #
+        # @param event [Symbol, nil] The event to count callbacks for, or nil for all events
+        # @return [Integer, Hash] Count for specific event, or hash of counts for all events
+        def callback_count(event = nil)
+          @callback_monitor.synchronize do
+            if event
+              @callbacks[event]&.size || 0
+            else
+              @callbacks.transform_values(&:size)
+            end
+          end
+        end
+
+        private
+
+        # Emits an event to all registered subscribers.
+        # Callbacks are executed in registration order (FIFO).
+        # Errors in callbacks are isolated - one failing callback doesn't prevent others from running.
+        #
+        # @param event [Symbol] The event to emit
+        # @param args [Array] Arguments to pass to each callback
+        # @return [void]
+        def emit(event, *args)
+          # Snapshot callbacks under lock (fast operation)
+          callbacks = @callback_monitor.synchronize { @callbacks[event]&.dup || [] }
+
+          # Execute callbacks outside lock (safe, non-blocking)
+          callbacks.each do |callback|
+            callback.call(*args)
+          rescue StandardError => e
+            handle_callback_error(event, callback, e)
+          end
+        end
+
+        # Hook for custom error handling when a callback raises an exception.
+        # Override this method in Chat to customize error behavior.
+        #
+        # @param event [Symbol] The event that was being emitted
+        # @param callback [Proc] The callback that raised the error
+        # @param error [StandardError] The error that was raised
+        # @return [void]
+        def handle_callback_error(event, _callback, error)
+          warn("[SwarmSDK] Callback error in #{event}: #{error.class} - #{error.message}")
+        end
+      end
+    end
+  end
+end

data/lib/swarm_sdk/agent/{chat → chat_helpers}/hook_integration.rb

@@ -2,7 +2,7 @@
 
 module SwarmSDK
   module Agent
-    class Chat < RubyLLM::Chat
+    module ChatHelpers
       # Integrates SwarmSDK's hook system with Agent::Chat
       #
       # Responsibilities:
@@ -71,40 +71,25 @@ module SwarmSDK
           @hook_agent_hooks[event].sort_by! { |cb| -cb.priority }
         end
 
-        # Override ask to trigger user_prompt hooks
+        # NOTE: The ask() method override has been removed.
         #
-        # This wraps the Agent::Chat#ask implementation to inject hooks AFTER
-        # system reminders are handled.
+        # In the new wrapper-based architecture, Agent::Chat#ask handles:
+        # 1. System reminder injection
+        # 2. User prompt hooks via trigger_user_prompt
+        # 3. Global semaphore acquisition
+        # 4. Delegation to RubyLLM::Chat
         #
-        # @param prompt [String] User prompt
-        # @param options [Hash] Additional options
-        # @return [RubyLLM::Message] LLM response
-        def ask(prompt, **options)
-          # Trigger user_prompt hook before sending to LLM (can halt or modify prompt)
-          if @hook_executor
-            hook_result = trigger_user_prompt(prompt)
-
-            # Check if hook halted execution
-            if hook_result[:halted]
-              # Return a halted message instead of calling LLM
-              return RubyLLM::Message.new(
-                role: :assistant,
-                content: hook_result[:halt_message],
-                model_id: model.id,
-              )
-            end
-
-            # Use modified prompt if hook provided one (stdout injection)
-            prompt = hook_result[:modified_prompt] if hook_result[:modified_prompt]
-          end
-
-          # Call original ask implementation (Agent::Chat handles system reminders)
-          super(prompt, **options)
-        end
+        # The hook integration is now done directly in Agent::Chat#ask rather than
+        # through module override, since there's no inheritance chain to call super on.
 
         # Override check_context_warnings to trigger our hook system
         #
         # This wraps the default context warning behavior to also trigger hooks.
+        # Unified implementation that:
+        # 1. Emits context_threshold_hit events (for snapshot reconstruction)
+        # 2. Optionally triggers automatic compression at 60% (if no custom handler)
+        # 3. Emits context_limit_warning events (backward compatibility)
+        # 4. Triggers user-defined context_warning hooks
         def check_context_warnings
           return unless respond_to?(:context_usage_percentage)
 
@@ -118,20 +103,40 @@ module SwarmSDK
             # Mark threshold as hit
             @agent_context.hit_warning_threshold?(threshold)
 
-            # Emit existing log event (for backward compatibility)
+            # Emit context_threshold_hit event (for snapshot reconstruction) - CRITICAL
+            LogStream.emit(
+              type: "context_threshold_hit",
+              agent: @agent_context.name,
+              threshold: threshold,
+              current_usage_percentage: current_percentage.round(2),
+            )
+
+            # Check if user has defined custom handler for context_warning
+            # Custom handlers take responsibility for managing context at this threshold
+            has_custom_handler = (@hook_agent_hooks[:context_warning] || []).any?
+
+            # Trigger automatic compression at 60% ONLY if no custom handler
+            compression_triggered = false
+            if threshold == Context::COMPRESSION_THRESHOLD && !has_custom_handler
+              compressed_count = apply_automatic_compression
+              compression_triggered = compressed_count > 0
+            end
+
+            # Emit legacy context_limit_warning for backwards compatibility
             LogStream.emit(
               type: "context_limit_warning",
               agent: @agent_context.name,
-              model: model.id,
+              model: model_id,
               threshold: "#{threshold}%",
               current_usage: "#{current_percentage}%",
               tokens_used: cumulative_total_tokens,
               tokens_remaining: tokens_remaining,
               context_limit: context_limit,
               metadata: @agent_context.metadata,
+              compression_triggered: compression_triggered,
             )
 
-            # Trigger hook system
+            # Trigger hook system (user-defined handlers)
             trigger_context_warning(threshold, current_percentage) if @hook_executor
           end
         end
@@ -186,9 +191,13 @@ module SwarmSDK
         def trigger_post_tool_use(result, tool_call:)
           return result unless @hook_executor
 
+          # Extract tracking digest for Read/MemoryRead tools
+          metadata_with_digest = extract_tool_tracking_digest(tool_call, result)
+
           context = build_hook_context(
             event: :post_tool_use,
             tool_result: wrap_tool_result(tool_call.id, tool_call.name, result),
+            metadata: metadata_with_digest,
           )
 
           agent_hooks = @hook_agent_hooks[:post_tool_use] || []
@@ -213,6 +222,45 @@ module SwarmSDK
 
         private
 
+        # Apply automatic message compression when context threshold is hit
+        #
+        # Called when context usage crosses 60% threshold and no custom handler exists.
+        # Compresses old tool results to save context window space while preserving accuracy.
+        #
+        # @return [Integer] Number of messages compressed (0 if compression not applied)
+        def apply_automatic_compression
+          return 0 unless respond_to?(:context_manager) && respond_to?(:messages)
+
+          # Calculate tokens before compression
+          tokens_before = cumulative_total_tokens
+
+          # Get compressed messages from ContextManager
+          compressed = context_manager.auto_compress_on_threshold(messages, keep_recent: 10)
+
+          # Count how many messages were actually compressed
+          messages_compressed = compressed.count do |msg|
+            msg.content.to_s.include?("[truncated for context management]")
+          end
+
+          # Replace messages using proper abstraction
+          replace_messages(compressed)
+
+          # Log compression event
+          LogStream.emit(
+            type: "context_compression",
+            agent: @agent_context.name,
+            total_messages: message_count,
+            messages_compressed: messages_compressed,
+            tokens_before: tokens_before,
+            current_usage: "#{context_usage_percentage}%",
+            compression_strategy: "progressive_tool_result_compression",
+            keep_recent: 10,
+            triggered_by: "auto_compression_threshold",
+          ) if LogStream.enabled?
+
+          messages_compressed
+        end
+
         # Trigger context_warning hooks
         #
         # Hooks have access to the chat instance via metadata[:chat]
@@ -251,13 +299,14 @@ module SwarmSDK
         # Can halt execution or append hook stdout to prompt.
         #
         # @param prompt [String] User's message/prompt
+        # @param source [String] Source of the prompt ("user" or "delegation")
         # @return [Hash] { halted: bool, halt_message: String, modified_prompt: String }
-        def trigger_user_prompt(prompt)
+        def trigger_user_prompt(prompt, source: "user")
           return { halted: false, modified_prompt: prompt } unless @hook_executor
 
-          # Filter out delegation tools from tools list
-          actual_tools = if respond_to?(:tools) && @agent_context
-            tools.keys.reject { |tool_name| @agent_context.delegation_tool?(tool_name.to_s) }
+          # Get tool names without delegation tools using proper abstraction
+          actual_tools = if respond_to?(:non_delegation_tool_names) && @agent_context
+            non_delegation_tool_names
           else
             []
           end
@@ -273,11 +322,12 @@ module SwarmSDK
             event: :user_prompt,
             metadata: {
               prompt: prompt,
-              message_count: messages.size,
-              model: model.id,
-              provider: model.provider,
+              message_count: message_count,
+              model: model_id,
+              provider: model_provider,
               tools: actual_tools,
               delegates_to: delegate_agents,
+              source: source,
               timestamp: Time.now.utc.iso8601,
             },
           )
@@ -335,6 +385,64 @@ module SwarmSDK
           )
         end
 
+        # Extract tracking digest for Read/MemoryRead tools
+        #
+        # Queries the appropriate tracker after tool execution to get the digest
+        # that was calculated and stored during the read operation.
+        #
+        # @param tool_call [RubyLLM::ToolCall] Tool call with arguments
+        # @param result [Object] Tool execution result (to check for errors)
+        # @return [Hash] Metadata hash with digest if applicable
+        def extract_tool_tracking_digest(tool_call, result)
+          # Only add digest for successful Read/MemoryRead tool calls
+          return {} if result.is_a?(StandardError)
+          return {} unless ["Read", "MemoryRead"].include?(tool_call.name)
+
+          # Extract path from arguments
+          path = case tool_call.name
+          when "Read"
+            tool_call.arguments[:file_path] || tool_call.arguments["file_path"]
+          when "MemoryRead"
+            tool_call.arguments[:file_path] || tool_call.arguments["file_path"]
+          end
+
+          return {} unless path
+
+          # Query tracker for digest
+          digest = case tool_call.name
+          when "Read"
+            Tools::Stores::ReadTracker.get_read_files(@agent_context.name)[File.expand_path(path)]
+          else
+            # Query registered plugins for digest (e.g., MemoryRead from SwarmMemory plugin)
+            query_plugin_for_digest(tool_call.name, path)
+          end
+
+          digest ? { read_digest: digest, read_path: path } : {}
+        end
+
+        # Query registered plugins for a tool result digest
+        #
+        # This allows plugins to provide digest tracking for their own tools
+        # (e.g., MemoryRead tracking in SwarmMemory plugin).
+        #
+        # @param tool_name [String] Name of the tool
+        # @param path [String] Path or identifier of the resource
+        # @return [String, nil] Digest from first plugin that responds, or nil
+        def query_plugin_for_digest(tool_name, path)
+          return unless @agent_context
+
+          PluginRegistry.all.each do |plugin|
+            digest = plugin.get_tool_result_digest(
+              agent_name: @agent_context.name,
+              tool_name: tool_name,
+              path: path,
+            )
+            return digest if digest
+          end
+
+          nil
+        end
+
         # Wrap a tool result in our Hooks::ToolResult value object
         #
         # @param tool_call_id [String] Tool call ID