Fix AI chat session job scoping to prevent wrong sessions when switching jobs

Author: elias-ba

CHANGELOG.md

@@ -42,6 +42,9 @@ and this project adheres to
 
 ### Fixed
 
+- Fixed AI chat session persistence when switching between jobs in workflow
+  editor [#3745](https://github.com/OpenFn/lightning/issues/3745)
+
 ## [2.14.12] - 2025-10-21
 
 ## [2.14.12-pre1] - 2025-10-21

lib/lightning_web/live/ai_assistant/component.ex

@@ -390,6 +390,8 @@ defmodule LightningWeb.AiAssistant.Component do
       :new ->
         socket
         |> delegate_to_handler(:on_session_close)
+        |> assign(:action, :new)
+        |> assign(:session, nil)
         |> assign_async([:all_sessions, :pagination_meta], fn ->
           case handler.list_sessions(assigns, sort_direction,
                  limit: @default_page_size
@@ -400,27 +402,80 @@ defmodule LightningWeb.AiAssistant.Component do
         end)
 
       :show ->
-        session = handler.get_session!(assigns)
-
-        message_loading =
-          Enum.any?(session.messages, fn msg ->
-            msg.role == :user && msg.status in [:pending, :processing]
-          end)
-
-        socket
-        |> assign(:session, session)
-        |> assign(
-          :pending_message,
-          if message_loading do
-            AsyncResult.loading()
-          else
-            AsyncResult.ok(nil)
-          end
-        )
-        |> delegate_to_handler(:on_session_open, [session])
+        try do
+          session = handler.get_session!(assigns)
+
+          message_loading =
+            Enum.any?(session.messages, fn msg ->
+              msg.role == :user && msg.status in [:pending, :processing]
+            end)
+
+          socket
+          |> assign(:session, session)
+          |> assign(:action, :show)
+          |> assign(
+            :pending_message,
+            if message_loading do
+              AsyncResult.loading()
+            else
+              AsyncResult.ok(nil)
+            end
+          )
+          |> delegate_to_handler(:on_session_open, [session])
+        rescue
+          Ecto.NoResultsError ->
+            # Session doesn't exist or doesn't belong to this job
+            # Fall back to showing the session list
+            socket
+            |> apply_action(:new)
+            |> assign(:action, :new)
+            |> assign(:session, nil)
+        end
     end
   end
 
+  defp load_and_show_session_list(socket, handler, assigns, sort_direction) do
+    socket
+    |> delegate_to_handler(:on_session_close)
+    |> assign(:action, :new)
+    |> assign(:session, nil)
+    |> assign_async([:all_sessions, :pagination_meta], fn ->
+      case handler.list_sessions(assigns, sort_direction,
+             limit: @default_page_size
+           ) do
+        %{sessions: sessions, pagination: pagination} ->
+          {:ok, %{all_sessions: sessions, pagination_meta: pagination}}
+      end
+    end)
+  end
+
+  defp load_and_show_session(socket, handler, assigns) do
+    session = handler.get_session!(assigns)
+
+    message_loading =
+      Enum.any?(session.messages, fn msg ->
+        msg.role == :user && msg.status in [:pending, :processing]
+      end)
+
+    socket
+    |> assign(:session, session)
+    |> assign(:action, :show)
+    |> assign(
+      :pending_message,
+      if message_loading do
+        AsyncResult.loading()
+      else
+        AsyncResult.ok(nil)
+      end
+    )
+    |> delegate_to_handler(:on_session_open, [session])
+  rescue
+    Ecto.NoResultsError ->
+      socket
+      |> apply_action(:new)
+      |> assign(:action, :new)
+      |> assign(:session, nil)
+  end
   defp save_message(socket, action, content) do
     result =
       case action do

lib/lightning_web/live/ai_assistant/modes/job_code.ex

@@ -100,7 +100,16 @@ defmodule LightningWeb.Live.AiAssistant.Modes.JobCode do
   @impl true
   @spec get_session!(assigns()) :: session()
   def get_session!(%{chat_session_id: session_id, selected_job: job} = assigns) do
-    AiAssistant.get_session!(session_id)
+    session = AiAssistant.get_session!(session_id)
+
+    # Validate that the session belongs to the selected job
+    if session.job_id != job.id do
+      raise Ecto.NoResultsError,
+        queryable: Lightning.AiAssistant.ChatSession,
+        message: "Chat session #{session_id} does not belong to job #{job.id}"
+    end
+
+    session
     |> AiAssistant.put_expression_and_adaptor(job.body, job.adaptor)
     |> maybe_add_run_logs(job, assigns[:follow_run])
   end

lib/lightning_web/live/workflow_live/edit.ex

@@ -1421,6 +1421,7 @@ defmodule LightningWeb.WorkflowLive.Edit do
        page_title: "",
        selected_edge: nil,
        selected_job: nil,
+       last_selected_job: nil,
        selected_run: nil,
        selected_trigger: nil,
        selection_mode: nil,
@@ -3584,10 +3585,19 @@ defmodule LightningWeb.WorkflowLive.Edit do
   end
 
   defp assign_chat_session_id(socket, params) do
+    job_chat_session_id =
+      if changed?(socket, :selected_job) &&
+           not is_nil(socket.assigns[:last_selected_job]) do
+        nil
+      else
+        params["j-chat"]
+      end
+
     socket
     |> assign(
       workflow_chat_session_id: params["w-chat"],
-      job_chat_session_id: params["j-chat"]
+      job_chat_session_id: job_chat_session_id,
+      last_selected_job: socket.assigns[:selected_job]
     )
   end
 

test/lightning/ai_assistant/ai_assistant_test.exs

@@ -2050,4 +2050,124 @@ defmodule Lightning.AiAssistantTest do
       assert AiAssistant.title_max_length() == 40
     end
   end
+
+  describe "list_sessions/2 - job scoping" do
+    test "only shows sessions for the currently selected job" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+      job_a = insert(:job, workflow: workflow, name: "Job A")
+      job_b = insert(:job, workflow: workflow, name: "Job B")
+
+      # Create sessions directly in the database (bypass AI processing)
+      session_a1 =
+        insert(:chat_session,
+          job: job_a,
+          user: user,
+          title: "Help with Job A - 1",
+          session_type: "job_code"
+        )
+
+      session_a2 =
+        insert(:chat_session,
+          job: job_a,
+          user: user,
+          title: "Help with Job A - 2",
+          session_type: "job_code"
+        )
+
+      _session_b =
+        insert(:chat_session,
+          job: job_b,
+          user: user,
+          title: "Help with Job B",
+          session_type: "job_code"
+        )
+
+      # List sessions for job_a
+      result = AiAssistant.list_sessions(job_a, :desc)
+
+      # Should only return job_a sessions
+      session_ids = Enum.map(result.sessions, & &1.id)
+      assert session_a1.id in session_ids
+      assert session_a2.id in session_ids
+      assert length(result.sessions) == 2
+    end
+
+    test "list_sessions for job_b doesn't include job_a sessions" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+      job_a = insert(:job, workflow: workflow, name: "Job A")
+      job_b = insert(:job, workflow: workflow, name: "Job B")
+
+      # Create sessions directly in the database
+      _session_a =
+        insert(:chat_session,
+          job: job_a,
+          user: user,
+          title: "Help with Job A",
+          session_type: "job_code"
+        )
+
+      session_b =
+        insert(:chat_session,
+          job: job_b,
+          user: user,
+          title: "Help with Job B",
+          session_type: "job_code"
+        )
+
+      # List sessions for job_b
+      result = AiAssistant.list_sessions(job_b, :desc)
+
+      # Should only return job_b session
+      session_ids = Enum.map(result.sessions, & &1.id)
+      assert session_b.id in session_ids
+      assert length(result.sessions) == 1
+    end
+
+    test "sessions from multiple jobs don't cross-contaminate" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+      job_a = insert(:job, workflow: workflow, name: "Job A")
+      job_b = insert(:job, workflow: workflow, name: "Job B")
+      job_c = insert(:job, workflow: workflow, name: "Job C")
+
+      # Create sessions for all jobs
+      session_a =
+        insert(:chat_session,
+          job: job_a,
+          user: user,
+          title: "Session A",
+          session_type: "job_code"
+        )
+
+      session_b =
+        insert(:chat_session,
+          job: job_b,
+          user: user,
+          title: "Session B",
+          session_type: "job_code"
+        )
+
+      session_c =
+        insert(:chat_session,
+          job: job_c,
+          user: user,
+          title: "Session C",
+          session_type: "job_code"
+        )
+
+      # Each job should only see its own session
+      result_a = AiAssistant.list_sessions(job_a, :desc)
+      result_b = AiAssistant.list_sessions(job_b, :desc)
+      result_c = AiAssistant.list_sessions(job_c, :desc)
+
+      assert [session_a.id] == Enum.map(result_a.sessions, & &1.id)
+      assert [session_b.id] == Enum.map(result_b.sessions, & &1.id)
+      assert [session_c.id] == Enum.map(result_c.sessions, & &1.id)
+    end
+  end
 end

test/lightning_web/ai_assistant/job_code_test.exs

@@ -1,6 +1,8 @@
 defmodule LightningWeb.AiAssistant.Modes.JobCodeTest do
   use Lightning.DataCase, async: true
 
+  import Lightning.Factories
+
   alias LightningWeb.Live.AiAssistant.Modes.JobCode
 
   describe "chat_input_disabled?/1" do
@@ -239,4 +241,91 @@ defmodule LightningWeb.AiAssistant.Modes.JobCodeTest do
       assert JobCode.chat_input_disabled?(assigns_ok) == false
     end
   end
+
+  describe "get_session!/1" do
+    test "successfully retrieves session when it belongs to the selected job" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+      job = insert(:job, workflow: workflow)
+
+      # Create session directly (bypass AI processing)
+      session =
+        insert(:chat_session,
+          job: job,
+          user: user,
+          session_type: "job_code"
+        )
+
+      assigns = %{
+        chat_session_id: session.id,
+        selected_job: job,
+        follow_run: nil
+      }
+
+      result = JobCode.get_session!(assigns)
+
+      assert result.id == session.id
+      assert result.job_id == job.id
+    end
+
+    test "raises Ecto.NoResultsError when session belongs to a different job" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+      job_a = insert(:job, workflow: workflow)
+      job_b = insert(:job, workflow: workflow)
+
+      # Create session for job_a
+      session =
+        insert(:chat_session,
+          job: job_a,
+          user: user,
+          session_type: "job_code"
+        )
+
+      # Try to get session with job_b selected
+      assigns = %{
+        chat_session_id: session.id,
+        selected_job: job_b,
+        follow_run: nil
+      }
+
+      assert_raise Ecto.NoResultsError, fn ->
+        JobCode.get_session!(assigns)
+      end
+    end
+
+    test "includes adaptor and expression from selected job" do
+      user = insert(:user)
+      project = insert(:project)
+      workflow = insert(:workflow, project: project)
+
+      job =
+        insert(:job,
+          workflow: workflow,
+          body: "fn(state => state)",
+          adaptor: "@openfn/language-http@1.0.0"
+        )
+
+      # Create session directly (bypass AI processing)
+      session =
+        insert(:chat_session,
+          job: job,
+          user: user,
+          session_type: "job_code"
+        )
+
+      assigns = %{
+        chat_session_id: session.id,
+        selected_job: job,
+        follow_run: nil
+      }
+
+      result = JobCode.get_session!(assigns)
+
+      assert result.expression == "fn(state => state)"
+      assert result.adaptor == "@openfn/language-http@1.0.0"
+    end
+  end
 end