llama_bot_rails 0.1.11 → 0.1.14

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ab8f2052c6695abc545b10879f23262d1c444011c28aa60e035fe94e12e62e35
-  data.tar.gz: 21532f0afe7f6bf345da06b89e0e472d6e219eba9a3da0a5387a09a2f3ddbc39
+  metadata.gz: 4a2a7cba5de6b95658edf2f1445ddd883bf0fa7b6f0a844fc7b65d75696bbe93
+  data.tar.gz: 471a4c54a3f26dd4d6b68c6d208d3c8e553cfcf2933436bad3e7b30af2d50296
 SHA512:
-  metadata.gz: a4400caf025fb253e8734003046a8256d6dcd6862ef6ddc9b45a3ac7a0a016749ef77eccf5ed9b17d5310ed1029a98af4873737765471d17aa665e0f60c7865c
-  data.tar.gz: 5ee28cfaa7c097dcce2fd95490c3721643abc0600cdbb3bd591bb255a2d21c27558810f43194206a10ae4ea0de7ee1ec71e9c5c0c62defb731238751fc8f1dd5
+  metadata.gz: 49afcfcec2697177cccc266cbdca5b7b0859feef23fb45d2626435c686368e0231d9cf0c66b894a3a0696dc50d54c9e285fc6b66905b42be752ac24126c82d97
+  data.tar.gz: ada3ecb0a74f2fef56122ffcc97889b7c877d670296a5f8a556d3afad685b53332f7be4306cfdf901f2b00744f8fd600c56041605aef02db666d9e51fe633a72

data/README.md

@@ -160,6 +160,23 @@ Rails.application.configure do
 end
 ```
 
+## **Principle of Least Priviledge**: Whitelisting Controller Actions
+
+Rather than giving the agent full access to your Rails app via the Rails console, You can whitelist controller actions instead.
+
+```
+class PagesController < ApplicationController
+  include LlamaBotRails::ControllerExtensions
+  include LlamaBotRails::AgentAuth
+  
+  # ─── Allow the agent to hit these actions ────────────────────────────────
+  llama_bot_allow :update #uses llama_bot_rails "authenticate_user_or_agent!" on top of your existing devise authentication.
+
+  skip_before_action :authenticate_user_or_agent!, only: [:show] #NOTE: You must change any skip_before_action callbacks that skip devise Authentication, to use :authenticate_user_or_agent!
+```
+
+When you include LlamaBotRails::AgentAuth, the gem aliases any authenticate_<scope>! filters to a unified guard. You can keep your old callbacks/skips for now, but you’ll see a deprecation warning—switch to authenticate_user_or_agent! at your convenience.
+
 ## 🧪 **What You Can Build**
 
 ### Developer Assistant

data/app/channels/llama_bot_rails/chat_channel.rb

@@ -47,11 +47,6 @@ module LlamaBotRails
         # Use a begin/rescue block to catch thread creation errors
       begin
 
-        @api_token = Rails.application.message_verifier(:llamabot_ws).generate(
-          { session_id: SecureRandom.uuid },
-          expires_in: 30.minutes
-        )
-
         @worker = Thread.new do
             Thread.current[:connection_id] = @connection_id
           Thread.current.abort_on_exception = true  # This will help surface errors
@@ -93,6 +88,16 @@ module LlamaBotRails
           end
         end
 
+        # Close the external WebSocket connection BEFORE stopping async tasks
+        if @external_ws_connection
+          begin
+            @external_ws_connection.close
+            Rails.logger.info "👋 [LlamaBot] Gracefully closed external WebSocket connection for: #{connection_id}"
+          rescue => e
+            Rails.logger.warn "❌ [LlamaBot] Could not close WebSocket connection: #{e.message}"
+          end
+        end
+
         # Clean up async tasks with better error handling
         begin
           @listener_task&.stop rescue nil
@@ -102,16 +107,6 @@ module LlamaBotRails
           Rails.logger.error "[LlamaBot] Error stopping async tasks: #{e.message}"
         end
         
-        # Clean up the connection
-        if @external_ws_connection
-          begin
-            @external_ws_connection.close
-            Rails.logger.info "[LlamaBot] Closed external WebSocket connection for: #{connection_id}"
-          rescue => e
-            Rails.logger.warn "[LlamaBot] Could not close WebSocket connection: #{e.message}"
-          end
-        end
-        
         # Force garbage collection in development/test environments to help clean up
         if !Rails.env.production?
           GC.start
@@ -126,7 +121,15 @@ module LlamaBotRails
     # through external websocket to FastAPI/Python backend.
     def receive(data)
       begin
-        #used to validate the message before it's sent to the backend.
+        #used to validate the message before it's sent to the llamabot-backend.
+
+        # Get the currently logged in user from the environment.
+        current_user = LlamaBotRails.current_user_resolver.call(connection.env)
+
+        @api_token = Rails.application.message_verifier(:llamabot_ws).generate(
+          { session_id: SecureRandom.uuid, user_id: current_user&.id},
+          expires_in: 30.minutes
+        )
 
         #This could be an example of how we might implement hooks & filters in the future.
         validate_message(data) #Placeholder for now, we are using this to mock errors being thrown. In the future, we can add actual validation logic.        
@@ -135,7 +138,7 @@ module LlamaBotRails
 
         builder = state_builder_class.new(
           params: data,
-          context: { api_token: @api_token }
+          context: { api_token: @api_token }.with_indifferent_access
         )
 
         # 2. Construct the LangGraph-ready state
@@ -148,6 +151,7 @@ module LlamaBotRails
         Rails.logger.info "[LlamaBot] Got message from Javascript LlamaBot Frontend: #{data.inspect}"
       rescue => e
         Rails.logger.error "[LlamaBot] Error in receive method: #{e.message}"
+        Rails.logger.error "[LlamaBot] Backtrace: #{e.backtrace.join("\n")}"
         send_message_to_frontend("error", e.message)
       end
     end
@@ -249,84 +253,82 @@ module LlamaBotRails
           # Wait for tasks to complete or connection to close
           [@listener_task, @keepalive_task].each(&:wait)
         rescue => e
-          Rails.logger.error "[LlamaBot] Failed to connect to external WebSocket for connection #{connection_id}: #{e.message}"
+          Rails.logger.error "❌ [LlamaBot] Failed to connect to external WebSocket for connection #{connection_id}: #{e.message}"
         ensure
           # Clean up tasks if they exist
           @listener_task&.stop
           @keepalive_task&.stop
-          @external_ws_connection&.close
+          if @external_ws_connection
+            @external_ws_connection.close
+            Rails.logger.info "👋 [LlamaBot] Cleaned up external WebSocket connection in ensure block"
+          end
         end
       end
     end
 
     # Listen for messages from the LlamaBot Backend
     def listen_to_external_websocket(connection)
-      while message = connection.read
+      begin
+        while message = connection.read
+          # Extract the actual message content
+          message_content = message.buffer if message.buffer
+          next unless message_content.present?
 
-        #Try to fix the ping/pong issue keepliave
-        # if message.type == :ping
-        
-        #   # respond with :pong
-        #   connection.write(Async::WebSocket::Messages::ControlFrame.new(:pong, frame.data))
-        #   connection.flush
-        #   next
-        # end
-        # Extract the actual message content
-        if message.buffer
-          message_content = message.buffer  # Use .data to get the message content
-        else
-          message_content = message.content
-        end
+          Rails.logger.info "[LlamaBot] Received from external WebSocket: #{message_content}"
 
-        Rails.logger.info "[LlamaBot] Received from external WebSocket: #{message_content}"
-        
-        begin
-          parsed_message = JSON.parse(message_content)
-          
-          if parsed_message["type"] != "pong"
-            # byebug
-          end
-          
-          case parsed_message["type"]
-          when "ai"
-            # Add any additional handling for write_code messages here
-            formatted_message = { message: {type: "ai", content: parsed_message['content'], base_message: parsed_message["base_message"]} }.to_json
-          when "tool"
-            # Add any additional handling for tool messages here
-            formatted_message = { message: {type: "tool", content: parsed_message['content'], base_message: parsed_message["base_message"]} }.to_json
-          when "error"
-            Rails.logger.error "[LlamaBot] ---------Received error message!----------"
-            response = parsed_message['content']
-            formatted_message = { message: message_content }.to_json
-            Rails.logger.error "[LlamaBot] ---------------------> Response: #{response}"
-            Rails.logger.error "[LlamaBot] ---------Completed error message!----------"
-          when "pong"
-            # Tell llamabot frontend that we've received a pong response, and we're still connected
-            formatted_message = { message: {type: "pong"} }.to_json
+          begin
+            parsed_message = JSON.parse(message_content)
+            
+            formatted_message = { message: {type: parsed_message["type"], content: parsed_message['content'], base_message: parsed_message["base_message"]} }.to_json
+            case parsed_message["type"]
+            when "error"
+              Rails.logger.error "[LlamaBot] ---------Received error message!----------"
+              response = parsed_message['content']
+              formatted_message = { message: message_content }.to_json
+              Rails.logger.error "[LlamaBot] ---------------------> Response: #{response}"
+              Rails.logger.error "[LlamaBot] ---------Completed error message!----------"
+            when "pong"
+              # Tell llamabot frontend that we've received a pong response, and we're still connected
+              formatted_message = { message: {type: "pong"} }.to_json
+            end
+            ActionCable.server.broadcast "chat_channel_#{params[:session_id]}", formatted_message
+          rescue JSON::ParserError => e
+            Rails.logger.error "[LlamaBot] Failed to parse message as JSON: #{e.message}"
+            # Continue to the next message without crashing the listener.
+            next
           end
-        rescue JSON::ParserError => e
-          Rails.logger.error "[LlamaBot] Failed to parse message as JSON: #{e.message}"
         end
-        ActionCable.server.broadcast "chat_channel_#{params[:session_id]}", formatted_message
+      rescue IOError, Errno::ECONNRESET => e
+        # This is a recoverable error. Log it and allow the task to end gracefully.
+        # The `ensure` block in `setup_external_websocket` will handle the cleanup.
+        Rails.logger.warn "❌ [LlamaBot] Connection lost while listening: #{e.message}. The connection will be closed."
       end
     end
 
     ###
     def send_keep_alive_pings(connection)
       loop do
-        ping_message = {
-          type: 'ping',
-          connection_id: @connection_id,
-          connection_state: !connection.closed? ? 'connected' : 'disconnected',
-          connection_class: connection.class.name
-        }.to_json
-        connection.write(ping_message)
-        connection.flush
-        Rails.logger.debug "[LlamaBot] Sent keep-alive ping: #{ping_message}"
+        # Stop the loop gracefully if the connection has already been closed.
+        break if connection.closed?
+        
+        begin
+          ping_message = {
+            type: 'ping',
+            connection_id: @connection_id,
+            connection_state: !connection.closed? ? 'connected' : 'disconnected',
+            connection_class: connection.class.name
+          }.to_json
+          connection.write(ping_message)
+          connection.flush
+          Rails.logger.debug "[LlamaBot] Sent keep-alive ping: #{ping_message}"
+        rescue IOError, Errno::ECONNRESET => e
+          Rails.logger.warn "❌ [LlamaBot] Could not send ping, connection likely closed: #{e.message}"
+          # Break the loop to allow the task to terminate gracefully.
+          break
+        end
+
         Async::Task.current.sleep(30)
       end
-    rescue => e
-      Rails.logger.error "[LlamaBot] Error in keep-alive ping: #{e.message} | Connection type: #{connection.class.name}"
     end
 
     # Send messages from the user to the LlamaBot Backend Socket

data/app/controllers/llama_bot_rails/agent_controller.rb

@@ -2,11 +2,25 @@ require 'llama_bot_rails/llama_bot'
 module LlamaBotRails
     class AgentController < ActionController::Base
         include ActionController::Live
+        include LlamaBotRails::AgentAuth
+        include LlamaBotRails::ControllerExtensions
+
+        llama_bot_allow :command
+        
         skip_before_action :verify_authenticity_token, only: [:command, :send_message]
-        before_action :authenticate_agent!, only: [:command]
+        # skip_before_action :authenticate_user_or_agent!, only: [:command]
+        
+        before_action :check_agent_authentication
 
         # POST /agent/command
         def command
+
+            unless LlamaBotRails.config.enable_console_tool
+                Rails.logger.warn("[[LlamaBot Debug]] Console tool is disabled")
+                render json: { error: "Console tool is disabled" }, status: :forbidden
+                return
+            end
+
             input = params[:command]
             
             # Capture both stdout and return value
@@ -16,18 +30,15 @@ module LlamaBotRails
             $stdout = output
             result = safety_eval(input)
             $stdout = STDOUT
-            
-            # If result is a string and output has content, prefer output
-            final_result = if output.string.present?
-                             output.string.strip
-                           else
-                             result
-                           end
-            
-            render json: { result: final_result }
+
+            # Return both output and result
+            render json: { 
+                output: output.string.strip, 
+                result: result 
+            }
         rescue => e
             $stdout = STDOUT  # Reset stdout on error
-            render json: { error: e.class.name, message: e.message }, status: :unprocessable_entity
+            render json: { error: "#{e.class.name}: #{e.message}" }
         end
 
         def index
@@ -78,15 +89,18 @@ module LlamaBotRails
             response.headers['Cache-Control'] = 'no-cache'
             response.headers['Connection']    = 'keep-alive'
 
+            # Get the currently logged in user from the environment.
+            current_user = LlamaBotRails.current_user_resolver.call(request.env)
+
             @api_token = Rails.application.message_verifier(:llamabot_ws).generate(
-                { session_id: SecureRandom.uuid },
+                { session_id: SecureRandom.uuid, user_id: current_user&.id},
                 expires_in: 30.minutes
-              )
+            )
 
             # 1. Instantiate the builder
             builder = state_builder_class.new(
                 params: params,
-                context: { api_token: @api_token }
+                context: { api_token: @api_token }.with_indifferent_access
             )
     
             # 2. Construct the LangGraph-ready state
@@ -131,25 +145,15 @@ module LlamaBotRails
         private 
 
         def safety_eval(input)
-            begin
-                # Change to Rails root directory for file operations
-                Dir.chdir(Rails.root) do
-                    # Create a safer evaluation context
-                    Rails.logger.info "[[LlamaBot]] Evaluating input: #{input}"
-                    binding.eval(input)
-                end
-            rescue => exception
-                Rails.logger.error "Error in safety_eval: #{exception.message}"
-                return exception.message
+            # Change to Rails root directory for file operations
+            Dir.chdir(Rails.root) do
+                # Create a safer evaluation context
+                Rails.logger.info "[[LlamaBot]] Evaluating input: #{input}"
+                binding.eval(input)
             end
-        end
-
-        def authenticate_agent!
-            auth_header = request.headers["Authorization"]
-            token = auth_header&.split("Bearer ")&.last  # Extract token after "Bearer "
-            @session_payload = Rails.application.message_verifier(:llamabot_ws).verify(token)
-        rescue ActiveSupport::MessageVerifier::InvalidSignature
-            head :unauthorized
+        rescue => exception
+            Rails.logger.error "Error in safety_eval: #{exception.message}"
+            raise exception
         end
 
         def state_builder_class

data/app/views/llama_bot_rails/agent/chat.html.erb

@@ -586,7 +586,7 @@
     let currentThreadId = null;
     let isSidebarCollapsed = false;
     let streamingTimeout = null;
-    const STREAMING_TIMEOUT_MS = 30000; // 30 seconds timeout
+    const STREAMING_TIMEOUT_MS = 3000000; // 3000 seconds timeout
 
     // Initialize the app
     document.addEventListener('DOMContentLoaded', function() {
@@ -968,6 +968,7 @@
                     }
                     
                   } catch (parseError) {
+                    addMessage(`Error: ${parseError} - Data: ${jsonData}`, 'error');
                     console.error('Error parsing SSE data:', parseError, 'Data:', jsonData);
                   }
                 }

data/app/views/llama_bot_rails/agent/chat_ws.html.erb

@@ -575,7 +575,7 @@ This deprecated and will be removed over time.
             <img src="https://service-jobs-images.s3.us-east-2.amazonaws.com/7rl98t1weu387r43il97h6ipk1l7" alt="LlamaBot Logo" class="logo">
             <div id="connectionStatusIconForLlamaBot" class="connection-status status-yellow"></div>
           </div>
-          <h1>LlamaBot Chat</h1>
+          <h1>Lenny the Llama</h1>
         </div>
         <button class="compose-button" onclick="startNewConversation()">
           <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
@@ -659,7 +659,10 @@ This deprecated and will be removed over time.
         },
         received(data) {
           const parsedData = JSON.parse(data).message;
+          console.log("LLM Response:", parsedData);
           switch (parsedData.type) {
+            case "AIMessageChunk":
+              addMessage(parsedData.content, parsedData.type, parsedData.base_message);
             case "ai":
               addMessage(parsedData.content, parsedData.type, parsedData.base_message);
               break;
@@ -966,6 +969,11 @@ This deprecated and will be removed over time.
       const messageDiv = document.createElement('div');
       messageDiv.className = `message ${sender}-message`;
       
+      if (sender == "AIMessageChunk"){
+        console.log("AIMessageChunk" + base_message);
+        messageDiv.innerHTML += text;
+      }
+      
       // Parse markdown for bot messages using Snarkdown, keep plain text for user messages
       if (sender === 'ai') { //Arghh. We're having issues with difference in formats between when we're streaming from updates mode, and when pulling state from checkpoint.
         if (text == ''){ //this is most likely a tool call.
@@ -1151,7 +1159,8 @@ This deprecated and will be removed over time.
 
         }
         else {
-          messageDiv.innerHTML = snarkdown(text); 
+          messageDiv.innerHTML = text;
+//          messageDiv.innerHTML = snarkdown(text); 
         }
       } else if (sender === 'tool') { //tool messages are not parsed as markdown
         if (base_message.name == 'run_rails_console_command') {

data/bin/bundle

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'bundler' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ENV['BUNDLER_VERSION'] = str
+
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('bundler', 'bundle', version)
+else
+gem "bundler", version
+load Gem.bin_path("bundler", "bundle", version)
+end

data/bin/bundler

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'bundler' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ENV['BUNDLER_VERSION'] = str
+
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('bundler', 'bundler', version)
+else
+gem "bundler", version
+load Gem.bin_path("bundler", "bundler", version)
+end

data/bin/byebug

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'byebug' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('byebug', 'byebug', version)
+else
+gem "byebug", version
+load Gem.bin_path("byebug", "byebug", version)
+end

data/bin/coderay

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'coderay' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('coderay', 'coderay', version)
+else
+gem "coderay", version
+load Gem.bin_path("coderay", "coderay", version)
+end

data/bin/erb

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'erb' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('erb', 'erb', version)
+else
+gem "erb", version
+load Gem.bin_path("erb", "erb", version)
+end

data/bin/faker

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'faker' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('faker', 'faker', version)
+else
+gem "faker", version
+load Gem.bin_path("faker", "faker", version)
+end

data/bin/htmldiff

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'diff-lcs' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('diff-lcs', 'htmldiff', version)
+else
+gem "diff-lcs", version
+load Gem.bin_path("diff-lcs", "htmldiff", version)
+end

data/bin/irb

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+#
+# This file was generated by RubyGems.
+#
+# The application 'irb' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require 'rubygems'
+
+Gem.use_gemdeps
+
+version = ">= 0.a"
+
+str = ARGV.first
+if str
+  str = str.b[/\A_(.*)_\z/, 1]
+  if str and Gem::Version.correct?(str)
+    version = str
+    ARGV.shift
+  end
+end
+
+if Gem.respond_to?(:activate_bin_path)
+load Gem.activate_bin_path('irb', 'irb', version)
+else
+gem "irb", version
+load Gem.bin_path("irb", "irb", version)
+end