From b6200478b79836d47c9767baebfa8b45e756fe06 Mon Sep 17 00:00:00 2001 From: "Elias W. BA" Date: Fri, 24 Oct 2025 01:09:11 +0000 Subject: [PATCH 1/5] Fix AI chat session job scoping to prevent wrong sessions when switching jobs --- CHANGELOG.md | 3 + .../live/ai_assistant/component.ex | 91 ++++-- .../live/ai_assistant/modes/job_code.ex | 11 +- lib/lightning_web/live/workflow_live/edit.ex | 12 +- .../ai_assistant/ai_assistant_test.exs | 120 ++++++++ .../ai_assistant/job_code_test.exs | 89 ++++++ .../live/ai_assistant_live_test.exs | 280 ++++++++++++++++++ 7 files changed, 586 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8de9df66f8..e9da3b0e65 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/lib/lightning_web/live/ai_assistant/component.ex b/lib/lightning_web/live/ai_assistant/component.ex index 1b8e5dfcac..c194ed9c15 100644 --- a/lib/lightning_web/live/ai_assistant/component.ex +++ b/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 diff --git a/lib/lightning_web/live/ai_assistant/modes/job_code.ex b/lib/lightning_web/live/ai_assistant/modes/job_code.ex index 8cc2fc2073..b61d1ccffa 100644 --- a/lib/lightning_web/live/ai_assistant/modes/job_code.ex +++ b/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 diff --git a/lib/lightning_web/live/workflow_live/edit.ex b/lib/lightning_web/live/workflow_live/edit.ex index 98818f78c3..98c86475af 100644 --- a/lib/lightning_web/live/workflow_live/edit.ex +++ b/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 diff --git a/test/lightning/ai_assistant/ai_assistant_test.exs b/test/lightning/ai_assistant/ai_assistant_test.exs index 97e79502b7..193d120ab9 100644 --- a/test/lightning/ai_assistant/ai_assistant_test.exs +++ b/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 diff --git a/test/lightning_web/ai_assistant/job_code_test.exs b/test/lightning_web/ai_assistant/job_code_test.exs index 909958d52d..34cba97af4 100644 --- a/test/lightning_web/ai_assistant/job_code_test.exs +++ b/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 diff --git a/test/lightning_web/live/ai_assistant_live_test.exs b/test/lightning_web/live/ai_assistant_live_test.exs index 1fec4578b2..f76e49a645 100644 --- a/test/lightning_web/live/ai_assistant_live_test.exs +++ b/test/lightning_web/live/ai_assistant_live_test.exs @@ -3450,4 +3450,284 @@ defmodule LightningWeb.AiAssistantLiveTest do project = insert(:project, project_users: [%{user: user, role: :owner}]) %{project: project} end + + defp create_workflow_with_jobs(%{project: project}) do + trigger = build(:trigger, type: :webhook) + + job_a = + build(:job, + body: ~s[fn(state => state)], + name: "Job A" + ) + + job_b = + build(:job, + body: ~s[fn(state => state)], + name: "Job B" + ) + + workflow = + build(:workflow, project: project) + |> with_job(job_a) + |> with_job(job_b) + |> with_trigger(trigger) + |> with_edge({trigger, job_a}, %{condition_type: :always}) + |> with_edge({job_a, job_b}, %{condition_type: :on_job_success}) + |> insert() + + {:ok, _snapshot} = Lightning.Workflows.Snapshot.create(workflow) + + workflow = Lightning.Repo.reload!(workflow) |> Lightning.Repo.preload(:jobs) + [job_a, job_b] = workflow.jobs + + %{workflow: workflow, job_a: job_a, job_b: job_b} + end + + describe "AI Assistant - Job Scoping: URL parameter clearing" do + setup [:create_workflow_with_jobs] + + test "clears j-chat parameter when switching from job A to job B", %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a, + job_b: job_b + } do + skip_disclaimer(user) + + session_a = + insert(:chat_session, + job: job_a, + user: user, + title: "Help with Job A", + session_type: "job_code" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, "j-chat": session_a.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_a.id}") + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_b.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_b.id}") + end + + test "preserves j-chat parameter when staying on same job", %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a + } do + skip_disclaimer(user) + + session_a = + insert(:chat_session, + job: job_a, + user: user, + title: "Help with Job A", + session_type: "job_code" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, "j-chat": session_a.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_a.id}") + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, "j-chat": session_a.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_a.id}") + end + + test "clears j-chat even when switching back to job with original session", + %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a, + job_b: job_b + } do + skip_disclaimer(user) + + 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" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, "j-chat": session_a.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_a.id}") + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_b.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_b.id}") + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id]}" + ) + + assert has_element?(view, "#job-pane-#{job_a.id}") + end + + test "handles nil selected_job gracefully", %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a + } do + skip_disclaimer(user) + + session_a = + insert(:chat_session, + job: job_a, + user: user, + title: "Help with Job A", + session_type: "job_code" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, "j-chat": session_a.id]}" + ) + + assert render(view) =~ "workflow" + end + end + + describe "AI Assistant - Job Scoping: Session validation and error handling" do + setup [:create_workflow_with_jobs] + + test "attempting to access wrong job's session via URL falls back gracefully", + %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a, + job_b: job_b + } do + skip_disclaimer(user) + + session_a = + insert(:chat_session, + job: job_a, + user: user, + title: "Help with Job A", + session_type: "job_code" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_b.id, m: "expand", "j-chat": session_a.id]}" + ) + + assert has_element?(view, "#job-#{job_b.id}-ai-assistant") + end + + test "nonexistent session ID doesn't crash the page", %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a + } do + skip_disclaimer(user) + + fake_session_id = Ecto.UUID.generate() + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, m: "expand", "j-chat": fake_session_id]}" + ) + + assert has_element?(view, "#job-#{job_a.id}-ai-assistant") + end + end + + describe "AI Assistant - Job Scoping: Session list filtering" do + setup [:create_workflow_with_jobs] + + test "job A's AI assistant only shows job A sessions", %{ + conn: conn, + project: project, + user: user, + workflow: workflow, + job_a: job_a, + job_b: job_b + } do + skip_disclaimer(user) + + 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" + ) + + {:ok, view, _html} = + live( + conn, + ~p"/projects/#{project.id}/w/#{workflow.id}?#{[v: workflow.lock_version, s: job_a.id, m: "expand"]}" + ) + + assert has_element?(view, "#job-#{job_a.id}-ai-assistant") + + result = Lightning.AiAssistant.list_sessions(job_a, :desc) + session_ids = Enum.map(result.sessions, & &1.id) + + assert session_a.id in session_ids + assert length(result.sessions) == 1 + end + end end From 25d9bc7e35e8cb1e70910e345635021c3a8dd839 Mon Sep 17 00:00:00 2001 From: "Elias W. BA" Date: Fri, 24 Oct 2025 01:28:46 +0000 Subject: [PATCH 2/5] Refactor session validation to use overloaded get_session! in AiAssistant context --- lib/lightning/ai_assistant/ai_assistant.ex | 18 ++++++++++++++++++ .../live/ai_assistant/modes/job_code.ex | 11 +---------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/lib/lightning/ai_assistant/ai_assistant.ex b/lib/lightning/ai_assistant/ai_assistant.ex index 8c9389f592..1f1095632c 100644 --- a/lib/lightning/ai_assistant/ai_assistant.ex +++ b/lib/lightning/ai_assistant/ai_assistant.ex @@ -274,6 +274,24 @@ defmodule Lightning.AiAssistant do end end + @doc """ + Gets a session scoped to a specific job. + + Returns the session only if it belongs to the given job. + Raises `Ecto.NoResultsError` if the session doesn't exist or doesn't belong to the job. + + ## Examples + + iex> get_session!(session_id, job) + %ChatSession{} + """ + def get_session!(session_id, %Job{id: job_id}) do + ChatSession + |> where([s], s.id == ^session_id and s.job_id == ^job_id) + |> Repo.one!() + |> Repo.preload([:user, messages: {session_messages_query(), :user}]) + end + def get_session(id) do case Repo.get(ChatSession, id) do nil -> diff --git a/lib/lightning_web/live/ai_assistant/modes/job_code.ex b/lib/lightning_web/live/ai_assistant/modes/job_code.ex index b61d1ccffa..22b76d4e37 100644 --- a/lib/lightning_web/live/ai_assistant/modes/job_code.ex +++ b/lib/lightning_web/live/ai_assistant/modes/job_code.ex @@ -100,16 +100,7 @@ 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 - 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.get_session!(session_id, job) |> AiAssistant.put_expression_and_adaptor(job.body, job.adaptor) |> maybe_add_run_logs(job, assigns[:follow_run]) end From 06b76a2976efb2b6163d1392060f1c8672c9ecf4 Mon Sep 17 00:00:00 2001 From: "Elias W. BA" Date: Fri, 24 Oct 2025 01:39:46 +0000 Subject: [PATCH 3/5] Extract session loading logic into helper functions for better readability --- .../live/ai_assistant/component.ex | 95 +++++++------------ 1 file changed, 32 insertions(+), 63 deletions(-) diff --git a/lib/lightning_web/live/ai_assistant/component.ex b/lib/lightning_web/live/ai_assistant/component.ex index c194ed9c15..caf810a600 100644 --- a/lib/lightning_web/live/ai_assistant/component.ex +++ b/lib/lightning_web/live/ai_assistant/component.ex @@ -388,49 +388,10 @@ defmodule LightningWeb.AiAssistant.Component do case action 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 - ) do - %{sessions: sessions, pagination: pagination} -> - {:ok, %{all_sessions: sessions, pagination_meta: pagination}} - end - end) + load_and_show_session_list(socket, handler, assigns, sort_direction) :show -> - 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 + load_and_show_session(socket, handler, assigns) end end @@ -450,32 +411,40 @@ defmodule LightningWeb.AiAssistant.Component do end defp load_and_show_session(socket, handler, assigns) do - session = handler.get_session!(assigns) + case fetch_session(handler, assigns) do + {:ok, session} -> + message_loading = + Enum.any?(session.messages, fn msg -> + msg.role == :user && msg.status in [:pending, :processing] + end) - 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]) - 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]) + :error -> + socket + |> apply_action(:new) + |> assign(:action, :new) + |> assign(:session, nil) + end + end + + defp fetch_session(handler, assigns) do + {:ok, handler.get_session!(assigns)} rescue - Ecto.NoResultsError -> - socket - |> apply_action(:new) - |> assign(:action, :new) - |> assign(:session, nil) + Ecto.NoResultsError -> :error end + defp save_message(socket, action, content) do result = case action do From 585e4b4ea8a80089769faea58d168db48ff13a84 Mon Sep 17 00:00:00 2001 From: Stuart Corbishley Date: Mon, 17 Nov 2025 14:43:17 +0200 Subject: [PATCH 4/5] Refactor get_session functions to follow Lightning conventions Refactor get_session implementation to use shared query functions and follow idiomatic Ecto patterns established elsewhere in the codebase: - Add session_query/2 helper for building queries with optional job scoping - Add preload_session/1 helper to centralize preloading logic - Change get_session/1 to return nil instead of {:error, :not_found} tuple - Add get_session/2 for job-scoped retrieval (returns nil if not found) - Add get_session!/2 for job-scoped retrieval (raises if not found) - Refactor get_session!/1 to use shared helpers This eliminates code duplication, uses single database queries with efficient preloading, and follows the pattern established by functions like get_project_user/2 and get_credential_body/2. Updated callers to handle nil returns instead of error tuples. --- lib/lightning/ai_assistant/ai_assistant.ex | 99 +++++++++++++------ .../live/ai_assistant/component.ex | 2 +- lib/lightning_web/live/workflow_live/edit.ex | 6 +- .../ai_assistant/ai_assistant_test.exs | 6 +- 4 files changed, 78 insertions(+), 35 deletions(-) diff --git a/lib/lightning/ai_assistant/ai_assistant.ex b/lib/lightning/ai_assistant/ai_assistant.ex index 1f1095632c..64b60eef42 100644 --- a/lib/lightning/ai_assistant/ai_assistant.ex +++ b/lib/lightning/ai_assistant/ai_assistant.ex @@ -239,6 +239,49 @@ defmodule Lightning.AiAssistant do |> handle_transaction_result() end + defp session_query(session_id, opts \\ []) do + job_id = Keyword.get(opts, :job_id) + + query = from(s in ChatSession, where: s.id == ^session_id) + + if job_id do + where(query, [s], s.job_id == ^job_id) + else + query + end + end + + defp preload_session(session) do + session = + Repo.preload(session, [:user, messages: {session_messages_query(), :user}]) + + if session.session_type == "workflow_template" do + Repo.preload(session, :project) + else + session + end + end + + @doc """ + Gets a session by ID. + + Returns the session or `nil` if not found. + + ## Examples + + iex> get_session(session_id) + %ChatSession{} + + iex> get_session("nonexistent") + nil + """ + def get_session(session_id) when is_binary(session_id) do + case session_query(session_id) |> Repo.one() do + nil -> nil + session -> preload_session(session) + end + end + @doc """ Retrieves a chat session by ID with all related data preloaded. @@ -261,47 +304,47 @@ defmodule Lightning.AiAssistant do `Ecto.NoResultsError` if no session exists with the given ID. """ - def get_session!(id) do - session = - ChatSession - |> Repo.get!(id) - |> Repo.preload([:user, messages: {session_messages_query(), :user}]) - - if session.session_type == "workflow_template" do - Repo.preload(session, :project) - else - session - end + def get_session!(session_id) when is_binary(session_id) do + session_query(session_id) + |> Repo.one!() + |> preload_session() end @doc """ Gets a session scoped to a specific job. - Returns the session only if it belongs to the given job. - Raises `Ecto.NoResultsError` if the session doesn't exist or doesn't belong to the job. + Returns the session or `nil` if not found or doesn't belong to the job. ## Examples - iex> get_session!(session_id, job) + iex> get_session(session_id, job) %ChatSession{} + + iex> get_session(session_id, different_job) + nil """ - def get_session!(session_id, %Job{id: job_id}) do - ChatSession - |> where([s], s.id == ^session_id and s.job_id == ^job_id) - |> Repo.one!() - |> Repo.preload([:user, messages: {session_messages_query(), :user}]) + def get_session(session_id, %Job{id: job_id}) do + case session_query(session_id, job_id: job_id) |> Repo.one() do + nil -> nil + session -> preload_session(session) + end end - def get_session(id) do - case Repo.get(ChatSession, id) do - nil -> - {:error, :not_found} + @doc """ + Gets a session scoped to a specific job, raising if not found. - session -> - {:ok, - session - |> Repo.preload([:user, messages: {session_messages_query(), :user}])} - end + ## Examples + + iex> get_session!(session_id, job) + %ChatSession{} + + iex> get_session!(session_id, different_job) + ** (Ecto.NoResultsError) + """ + def get_session!(session_id, %Job{id: job_id}) do + session_query(session_id, job_id: job_id) + |> Repo.one!() + |> preload_session() end defp session_messages_query do diff --git a/lib/lightning_web/live/ai_assistant/component.ex b/lib/lightning_web/live/ai_assistant/component.ex index caf810a600..a5150bd3bf 100644 --- a/lib/lightning_web/live/ai_assistant/component.ex +++ b/lib/lightning_web/live/ai_assistant/component.ex @@ -328,7 +328,7 @@ defmodule LightningWeb.AiAssistant.Component do |> AiAssistant.retry_message() |> case do {:ok, {_message, _oban_job}} -> - {:ok, session} = AiAssistant.get_session(socket.assigns.session.id) + session = AiAssistant.get_session!(socket.assigns.session.id) {:noreply, socket diff --git a/lib/lightning_web/live/workflow_live/edit.ex b/lib/lightning_web/live/workflow_live/edit.ex index 98c86475af..1c76e7896a 100644 --- a/lib/lightning_web/live/workflow_live/edit.ex +++ b/lib/lightning_web/live/workflow_live/edit.ex @@ -3721,12 +3721,12 @@ defmodule LightningWeb.WorkflowLive.Edit do }) when is_binary(chat_id) do case Lightning.AiAssistant.get_session(chat_id) do - {:ok, session} -> + %Lightning.AiAssistant.ChatSession{} = session -> Lightning.AiAssistant.associate_workflow(session, workflow) - {:error, reason} -> + nil -> Logger.warning( - "Failed to associate workflow with chat session #{chat_id}: #{inspect(reason)}" + "Failed to associate workflow with chat session #{chat_id}: not found" ) end end diff --git a/test/lightning/ai_assistant/ai_assistant_test.exs b/test/lightning/ai_assistant/ai_assistant_test.exs index 193d120ab9..73ac5d7989 100644 --- a/test/lightning/ai_assistant/ai_assistant_test.exs +++ b/test/lightning/ai_assistant/ai_assistant_test.exs @@ -938,12 +938,12 @@ defmodule Lightning.AiAssistantTest do } do session = insert(:chat_session, user: user, job: job) - assert {:ok, retrieved_session} = AiAssistant.get_session(session.id) + assert retrieved_session = AiAssistant.get_session(session.id) assert retrieved_session.id == session.id end - test "returns {:error, :not_found} when not found" do - assert {:error, :not_found} = AiAssistant.get_session(Ecto.UUID.generate()) + test "returns nil when not found" do + assert nil == AiAssistant.get_session(Ecto.UUID.generate()) end end From beaf152382f6efe555cf6362e879a2ee3fea23bb Mon Sep 17 00:00:00 2001 From: Stuart Corbishley Date: Mon, 17 Nov 2025 14:57:48 +0200 Subject: [PATCH 5/5] Move CHANGELOG line to top --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bd2e40b88..109460199d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,8 @@ 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) - Fix ghost edges blocking saves and breaking autolayout when deleting jobs in collaborative editor [#3983](https://github.com/OpenFn/lightning/issues/3983) - Fix tooltip styling inconsistencies in collaborative editor @@ -252,8 +254,6 @@ 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) - Fixed credential preservation during sandbox workflow merge - credentials are now correctly maintained when merging sandboxes back to parent projects [#3782](https://github.com/OpenFn/lightning/issues/3782)