fix: AI assistant not showing logs when run selected mid-session

CHANGELOG.md

@@ -26,6 +26,8 @@ and this project adheres to
 
 ### Fixed
 
+- AI assistant now correctly displays logs when run is selected mid-session
+  [#4380](https://github.com/OpenFn/lightning/issues/4380)
 - Fix duplicate "Log in" heading on login page
   [#4459](https://github.com/OpenFn/lightning/issues/4459)
 - Editing an OAuth credential from the workflow canvas incorrectly showed an

lib/lightning/ai_assistant/message_processor.ex

@@ -92,21 +92,35 @@ defmodule Lightning.AiAssistant.MessageProcessor do
   defp update_session_with_job_context(session, message) do
     message_meta = message.meta || %{}
 
-    cond do
-      message.job_id ->
-        %{session | job_id: message.job_id}
+    session =
+      cond do
+        message.job_id ->
+          %{session | job_id: message.job_id}
+
+        Map.has_key?(message_meta, "unsaved_job") ->
+          updated_meta =
+            Map.put(
+              session.meta || %{},
+              "unsaved_job",
+              message_meta["unsaved_job"]
+            )
+
+          %{session | meta: updated_meta}
+
+        true ->
+          session
+      end
 
-      Map.has_key?(message_meta, "unsaved_job") ->
+    # Propagate follow_run_id from message meta to in-memory session meta
+    # so that maybe_add_run_logs can find it during enrichment
+    case message_meta do
+      %{"follow_run_id" => run_id} when not is_nil(run_id) ->
         updated_meta =
-          Map.put(
-            session.meta || %{},
-            "unsaved_job",
-            message_meta["unsaved_job"]
-          )
+          Map.put(session.meta || %{}, "follow_run_id", run_id)
 
         %{session | meta: updated_meta}
 
-      true ->
+      _ ->
         session
     end
   end

lib/lightning_web/channels/ai_assistant_channel.ex

@@ -631,6 +631,14 @@ defmodule LightningWeb.AiAssistantChannel do
     []
   end
 
+  defp maybe_put_follow_run_id_in_meta(attrs, %{"follow_run_id" => run_id})
+       when not is_nil(run_id) do
+    existing_meta = Map.get(attrs, :meta, %{})
+    Map.put(attrs, :meta, Map.put(existing_meta, "follow_run_id", run_id))
+  end
+
+  defp maybe_put_follow_run_id_in_meta(attrs, _params), do: attrs
+
   defp build_message_options(params) do
     %{
       "code" => params["attach_code"] == true,
@@ -838,7 +846,10 @@ defmodule LightningWeb.AiAssistantChannel do
        ) do
     case may_get_job(params["job_id"]) do
       {:ok, job} ->
-        message_attrs = build_message_attrs(user, job, content, limit_result)
+        message_attrs =
+          build_message_attrs(user, job, content, limit_result)
+          |> maybe_put_follow_run_id_in_meta(params)
+
         opts = extract_message_options(params)
 
         case AiAssistant.save_message(session, message_attrs, opts) do
@@ -896,6 +907,7 @@ defmodule LightningWeb.AiAssistantChannel do
     message_attrs =
       build_message_attrs(user, nil, content, limit_result)
       |> Map.put(:meta, %{"unsaved_job" => unsaved_job_data})
+      |> maybe_put_follow_run_id_in_meta(params)
 
     opts = extract_message_options(params)
 

test/lightning/ai_assistant/ai_assistant_test.exs

@@ -1281,6 +1281,54 @@ defmodule Lightning.AiAssistantTest do
                  "@openfn/language-http@latest"
                )
     end
+
+    test "fetches logs when follow_run_id is added mid-session", %{
+      user: user,
+      workflow: %{jobs: [job | _]} = workflow
+    } do
+      # Simulate the bug scenario: session created without follow_run_id,
+      # then user selects a run later during the session
+      work_order = insert(:workorder, workflow: workflow)
+
+      run =
+        insert(:run,
+          work_order: work_order,
+          dataclip: build(:dataclip),
+          starting_job: job
+        )
+
+      step = insert(:step, job: job)
+      insert(:run_step, run: run, step: step)
+
+      insert(:log_line,
+        step: step,
+        run: run,
+        message: "Debug log from mid-session run",
+        timestamp: ~U[2024-01-01 10:00:00Z]
+      )
+
+      # Session initially created without follow_run_id
+      session =
+        insert(:chat_session,
+          user: user,
+          job: job,
+          meta: %{}
+        )
+
+      # Verify no logs without follow_run_id
+      enriched_before = AiAssistant.enrich_session_with_job_context(session)
+      assert is_nil(enriched_before.logs)
+
+      # Now simulate follow_run_id being added via message processing
+      # (this is what the message processor does after the fix)
+      session_with_run = %{session | meta: %{"follow_run_id" => run.id}}
+
+      enriched_after =
+        AiAssistant.enrich_session_with_job_context(session_with_run)
+
+      # Logs should now be present
+      assert enriched_after.logs =~ "Debug log from mid-session run"
+    end
   end
 
   describe "retry_message/1" do

test/lightning/ai_assistant/message_processor_test.exs

@@ -197,5 +197,64 @@ defmodule Lightning.AiAssistant.MessageProcessorTest do
       assert is_nil(assistant_message.job_id)
       refute Map.has_key?(assistant_message.meta || %{}, "from_unsaved_job")
     end
+
+    @tag :capture_log
+    test "processes message successfully when follow_run_id is in message.meta",
+         %{
+           user: user,
+           project: project
+         } do
+      workflow = insert(:workflow, project: project)
+      job = insert(:job, workflow: workflow)
+
+      # Create a run for the job
+      work_order = insert(:workorder, workflow: workflow)
+
+      run =
+        insert(:run,
+          work_order: work_order,
+          dataclip: build(:dataclip),
+          starting_job: job
+        )
+
+      # Create session without follow_run_id (user hasn't selected a run yet)
+      session =
+        insert(:chat_session,
+          user: user,
+          session_type: "job_code",
+          project: project,
+          job_id: job.id,
+          meta: %{}
+        )
+
+      # User later selects a run mid-session, sending follow_run_id in message params
+      {:ok, updated_session} =
+        AiAssistant.save_message(
+          session,
+          %{
+            role: :user,
+            content: "help me debug these logs",
+            user: user,
+            job: job,
+            meta: %{"follow_run_id" => run.id}
+          },
+          []
+        )
+
+      user_message = Enum.find(updated_session.messages, &(&1.role == :user))
+      assert user_message.meta["follow_run_id"] == run.id
+
+      # Process the message - update_session_with_job_context should use
+      # follow_run_id from message.meta for enrichment
+      assert :ok =
+               perform_job(MessageProcessor, %{"message_id" => user_message.id})
+
+      reloaded = AiAssistant.get_session!(session.id)
+      assistant_message = Enum.find(reloaded.messages, &(&1.role == :assistant))
+
+      # Verify assistant message was created successfully (proving enrichment worked)
+      assert assistant_message != nil
+      assert assistant_message.job_id == job.id
+    end
   end
 end

test/lightning_web/channels/ai_assistant_channel_test.exs

@@ -640,6 +640,61 @@ defmodule LightningWeb.AiAssistantChannelTest do
         assert user_msg.meta["unsaved_job"]["id"] == unknown_job_id
       end)
     end
+
+    @tag :capture_log
+    test "stores follow_run_id in message.meta when provided",
+         %{
+           socket: socket,
+           job: job,
+           user: user,
+           workflow: workflow,
+           project: project
+         } do
+      # Create a run for the job
+      dataclip = insert(:dataclip, project: project)
+      work_order = insert(:workorder, workflow: workflow)
+
+      run =
+        insert(:run,
+          work_order: work_order,
+          starting_job: job,
+          dataclip: dataclip
+        )
+
+      {:ok, session} =
+        AiAssistant.create_session(job, user, "Initial message", [])
+
+      {:ok, _, socket} =
+        subscribe_and_join(
+          socket,
+          AiAssistantChannel,
+          "ai_assistant:job_code:#{session.id}",
+          %{}
+        )
+
+      with_testing_mode(:manual, fn ->
+        # Simulate user selecting a run and checking "Send logs"
+        ref =
+          push(socket, "new_message", %{
+            "content" => "Help me debug these logs",
+            "follow_run_id" => run.id,
+            "attach_logs" => true
+          })
+
+        assert_reply ref, :ok, %{message: message}
+        assert message.role == "user"
+
+        # Verify follow_run_id was stored in message.meta
+        reloaded = AiAssistant.get_session!(session.id)
+
+        user_msg =
+          Enum.find(reloaded.messages, fn msg ->
+            msg.role == :user && msg.content == "Help me debug these logs"
+          end)
+
+        assert user_msg.meta["follow_run_id"] == run.id
+      end)
+    end
   end
 
   describe "handle_in mark_disclaimer_read" do