Refactor session validation to use overloaded get_session! in AiAssistant context

Author: elias-ba

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 ->

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