diff --git a/CHANGELOG.md b/CHANGELOG.md index 348716952b..d5e221262a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/lib/lightning/ai_assistant/message_processor.ex b/lib/lightning/ai_assistant/message_processor.ex index e45729a476..61ba48fc8c 100644 --- a/lib/lightning/ai_assistant/message_processor.ex +++ b/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 diff --git a/lib/lightning_web/channels/ai_assistant_channel.ex b/lib/lightning_web/channels/ai_assistant_channel.ex index ba676b46bb..1dd621280b 100644 --- a/lib/lightning_web/channels/ai_assistant_channel.ex +++ b/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) diff --git a/test/lightning/ai_assistant/ai_assistant_test.exs b/test/lightning/ai_assistant/ai_assistant_test.exs index 5ef5ae1d5a..2c549ab0b6 100644 --- a/test/lightning/ai_assistant/ai_assistant_test.exs +++ b/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 diff --git a/test/lightning/ai_assistant/message_processor_test.exs b/test/lightning/ai_assistant/message_processor_test.exs index 82c31d46b0..ef18e68cd5 100644 --- a/test/lightning/ai_assistant/message_processor_test.exs +++ b/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 diff --git a/test/lightning_web/channels/ai_assistant_channel_test.exs b/test/lightning_web/channels/ai_assistant_channel_test.exs index 35f5bf69bd..f45e8e63ac 100644 --- a/test/lightning_web/channels/ai_assistant_channel_test.exs +++ b/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