optional: nicer apollo errors for lightning users

CHANGELOG.md

@@ -47,6 +47,8 @@ and this project adheres to
   by adding `Run.active_states/0`, `WorkOrder.states/0`, and
   `WorkOrder.active_states/0` and replacing all hardcoded state lists across the
   codebase [#4589](https://github.com/OpenFn/lightning/issues/4589)
+- Add nicer error messages in the AI chat for when tokens don't work
+  [PR#4879](https://github.com/OpenFn/lightning/pull/4879)
 
 ### Fixed
 

assets/js/collaborative-editor/components/MessageList.tsx

@@ -652,7 +652,8 @@ export function MessageList({
                     >
                       <span className="hero-exclamation-circle h-4 w-4 text-red-600 flex-shrink-0" />
                       <span className="text-sm text-red-700 flex-1">
-                        Failed to send message. Please try again.
+                        {message.error ||
+                          'Failed to send message. Please try again.'}
                       </span>
                       {onRetryMessage && (
                         <button
@@ -748,7 +749,7 @@ export function MessageList({
                     >
                       <span className="hero-exclamation-circle h-3.5 w-3.5 text-red-600" />
                       <span className="text-xs text-red-700 flex-1">
-                        Failed to send
+                        {message.error || 'Failed to send'}
                       </span>
                       {onRetryMessage && (
                         <button

assets/js/collaborative-editor/lib/AIChannelRegistry.ts

@@ -678,8 +678,13 @@ export class AIChannelRegistry {
       const typedPayload = payload as {
         message_id: string;
         status: MessageStatus;
+        error?: string;
       };
-      this.store._updateMessageStatus(typedPayload.message_id, 'error');
+      this.store._updateMessageStatus(
+        typedPayload.message_id,
+        'error',
+        typedPayload.error
+      );
       this.store._setProcessingState(false);
     };
 

assets/js/collaborative-editor/stores/createAIAssistantStore.ts

@@ -444,12 +444,20 @@ export const createAIAssistantStore = (): AIAssistantStore => {
    * Update message status
    * @internal Called by useAIAssistantChannel hook
    */
-  const _updateMessageStatus = (messageId: string, status: MessageStatus) => {
+  const _updateMessageStatus = (
+    messageId: string,
+    status: MessageStatus,
+    errorMessage?: string
+  ) => {
     state = produce(state, draft => {
       const message = draft.messages.find(m => m.id === messageId);
       if (message) {
         message.status = status;
 
+        if (status === 'error' && errorMessage) {
+          message.error = errorMessage;
+        }
+
         if (status === 'success' || status === 'error') {
           draft.isLoading = false;
         }

assets/js/collaborative-editor/types/ai-assistant.ts

@@ -54,6 +54,8 @@ export interface Message {
   user_id?: string;
   user?: MessageUser | null;
   job_id?: string;
+  /** Human-readable error surfaced when status is "error" (e.g. Apollo 401). */
+  error?: string;
 }
 
 /**
@@ -190,7 +192,11 @@ export interface AIAssistantStore {
   _clearSessionList: () => void;
   _prependSession: (session: SessionSummary) => void;
   _addMessage: (message: Message) => void;
-  _updateMessageStatus: (messageId: string, status: MessageStatus) => void;
+  _updateMessageStatus: (
+    messageId: string,
+    status: MessageStatus,
+    errorMessage?: string
+  ) => void;
   _setSessionList: (response: SessionListResponse) => void;
   _appendSessionList: (response: SessionListResponse) => void;
   _initializeContext: (

lib/lightning/ai_assistant/ai_assistant.ex

@@ -1424,6 +1424,11 @@ defmodule Lightning.AiAssistant do
 
   defp handle_error_response(error_response, session) do
     case error_response do
+      {:ok, %Tesla.Env{status: 401}} ->
+        Logger.error("AI query unauthorized (401) for session #{session.id}")
+
+        {:error, "Token invalid or limit reached. Contact your administrator."}
+
       {:ok, %Tesla.Env{status: status, body: body}}
       when status not in @success_status_range ->
         error_message = body["message"]

lib/lightning/ai_assistant/message_processor.ex

@@ -184,7 +184,7 @@ defmodule Lightning.AiAssistant.MessageProcessor do
 
       {:error, error_message} ->
         {:ok, _updated_session, _updated_message} =
-          update_message_status(message, :error)
+          update_message_status(message, :error, error_message)
 
         {:error, error_message}
     end
@@ -290,11 +290,15 @@ defmodule Lightning.AiAssistant.MessageProcessor do
   """
   @spec update_message_status(
           ChatMessage.t(),
-          atom()
+          atom(),
+          String.t() | nil
         ) ::
           {:ok, ChatSession.t(), ChatMessage.t()}
-  def update_message_status(message, status) do
-    changes = build_status_changes(status)
+  def update_message_status(message, status, error_message \\ nil) do
+    changes =
+      status
+      |> build_status_changes()
+      |> maybe_put_error_message(message, error_message)
 
     updated_message =
       message
@@ -331,6 +335,19 @@ defmodule Lightning.AiAssistant.MessageProcessor do
     }
   end
 
+  # Persists the error string onto the message meta so it can be surfaced to
+  # the frontend (live via the message_error broadcast, and on reload via
+  # format_message) instead of a generic "Failed to send".
+  @spec maybe_put_error_message(map(), ChatMessage.t(), String.t() | nil) ::
+          map()
+  defp maybe_put_error_message(changes, message, error_message)
+       when is_binary(error_message) do
+    meta = Map.put(message.meta || %{}, "error_message", error_message)
+    Map.put(changes, :meta, meta)
+  end
+
+  defp maybe_put_error_message(changes, _message, _error_message), do: changes
+
   @doc false
   @spec workflow_code_from_session(AiAssistant.ChatSession.t()) ::
           String.t() | nil

lib/lightning_web/channels/ai_assistant_channel.ex

@@ -291,7 +291,8 @@ defmodule LightningWeb.AiAssistantChannel do
           if user_message do
             broadcast(socket, "message_error", %{
               message_id: user_message.id,
-              status: "error"
+              status: "error",
+              error: (user_message.meta || %{})["error_message"]
             })
           end
 
@@ -742,7 +743,8 @@ defmodule LightningWeb.AiAssistantChannel do
       inserted_at: message.inserted_at,
       user_id: message.user_id,
       user: format_user(message.user),
-      job_id: job_id
+      job_id: job_id,
+      error: (message.meta || %{})["error_message"]
     }
   end
 

test/lightning/ai_assistant/message_processor_test.exs

@@ -386,4 +386,153 @@ defmodule Lightning.AiAssistant.MessageProcessorTest do
       assert assistant_msg.content == "Workflow response"
     end
   end
+
+  describe "error message surfacing" do
+    test "persists Apollo 401 errors to message meta as a friendly message", %{
+      user: user,
+      project: project
+    } do
+      workflow = insert(:workflow, project: project)
+
+      session =
+        insert(:chat_session,
+          user: user,
+          session_type: "workflow_template",
+          project: project,
+          workflow: workflow,
+          job_id: nil,
+          meta: %{}
+        )
+
+      {:ok, updated_session} =
+        AiAssistant.save_message(
+          session,
+          %{role: :user, content: "generate a workflow", user: user},
+          []
+        )
+
+      user_message = Enum.find(updated_session.messages, &(&1.role == :user))
+
+      # Apollo rejects the request with a 401 (bad/expired token or limit hit)
+      Mox.stub(Lightning.Tesla.Mock, :call, fn
+        %{method: :get}, _opts -> {:ok, %Tesla.Env{status: 200}}
+        %{method: :post}, _opts -> {:ok, %Tesla.Env{status: 401, body: %{}}}
+      end)
+
+      assert :ok =
+               perform_job(MessageProcessor, %{"message_id" => user_message.id})
+
+      errored =
+        AiAssistant.get_session!(session.id).messages
+        |> Enum.find(&(&1.id == user_message.id))
+
+      assert errored.status == :error
+
+      assert errored.meta["error_message"] ==
+               "Token invalid or limit reached. Contact your administrator."
+    end
+
+    test "persists the Apollo body message to meta on other failures", %{
+      user: user,
+      project: project
+    } do
+      workflow = insert(:workflow, project: project)
+
+      session =
+        insert(:chat_session,
+          user: user,
+          session_type: "workflow_template",
+          project: project,
+          workflow: workflow,
+          job_id: nil,
+          meta: %{}
+        )
+
+      {:ok, updated_session} =
+        AiAssistant.save_message(
+          session,
+          %{role: :user, content: "generate a workflow", user: user},
+          []
+        )
+
+      user_message = Enum.find(updated_session.messages, &(&1.role == :user))
+
+      Mox.stub(Lightning.Tesla.Mock, :call, fn
+        %{method: :get}, _opts ->
+          {:ok, %Tesla.Env{status: 200}}
+
+        %{method: :post}, _opts ->
+          {:ok,
+           %Tesla.Env{status: 500, body: %{"message" => "Apollo is on fire"}}}
+      end)
+
+      assert :ok =
+               perform_job(MessageProcessor, %{"message_id" => user_message.id})
+
+      errored =
+        AiAssistant.get_session!(session.id).messages
+        |> Enum.find(&(&1.id == user_message.id))
+
+      assert errored.status == :error
+      assert errored.meta["error_message"] == "Apollo is on fire"
+    end
+
+    test "update_message_status/3 merges the error message into meta without clobbering existing keys",
+         %{user: user, project: project} do
+      workflow = insert(:workflow, project: project)
+
+      session =
+        insert(:chat_session,
+          user: user,
+          session_type: "workflow_template",
+          project: project,
+          workflow: workflow
+        )
+
+      message =
+        insert(:chat_message,
+          role: :user,
+          status: :processing,
+          user: user,
+          chat_session: session,
+          meta: %{"follow_run_id" => "keep-me"}
+        )
+
+      {:ok, _session, updated} =
+        MessageProcessor.update_message_status(message, :error, "boom")
+
+      assert updated.status == :error
+      assert updated.meta["error_message"] == "boom"
+      assert updated.meta["follow_run_id"] == "keep-me"
+    end
+
+    test "update_message_status/2 leaves meta untouched when no error message is given",
+         %{user: user, project: project} do
+      workflow = insert(:workflow, project: project)
+
+      session =
+        insert(:chat_session,
+          user: user,
+          session_type: "workflow_template",
+          project: project,
+          workflow: workflow
+        )
+
+      message =
+        insert(:chat_message,
+          role: :user,
+          status: :processing,
+          user: user,
+          chat_session: session,
+          meta: %{"follow_run_id" => "keep-me"}
+        )
+
+      {:ok, _session, updated} =
+        MessageProcessor.update_message_status(message, :error)
+
+      assert updated.status == :error
+      refute Map.has_key?(updated.meta, "error_message")
+      assert updated.meta["follow_run_id"] == "keep-me"
+    end
+  end
 end

test/lightning_web/channels/ai_assistant_channel_test.exs

@@ -2688,6 +2688,63 @@ defmodule LightningWeb.AiAssistantChannelTest do
       }
     end
 
+    test "broadcasts message_error including the persisted error message", %{
+      socket: socket,
+      job: job,
+      user: user
+    } do
+      {:ok, session} =
+        AiAssistant.create_session(job, user, "Initial message", [])
+
+      {:ok, _, socket} =
+        subscribe_and_join(
+          socket,
+          AiAssistantChannel,
+          "ai_assistant:job_code:#{session.id}",
+          %{}
+        )
+
+      {:ok, updated_session} =
+        AiAssistant.save_message(
+          session,
+          %{role: :user, content: "Test message", user: user},
+          []
+        )
+
+      message = List.last(updated_session.messages)
+      message_id = message.id
+
+      # Persist an error message onto the user message meta, mirroring what
+      # MessageProcessor does when Apollo returns an error.
+      message
+      |> Ecto.Changeset.change(
+        meta: %{
+          "error_message" =>
+            "Token invalid or limit reached. Contact your administrator."
+        }
+      )
+      |> Lightning.Repo.update!()
+
+      session_with_error = AiAssistant.get_session!(session.id)
+
+      send(
+        socket.channel_pid,
+        {:ai_assistant, :message_status_changed,
+         %{
+           status: {:error, session_with_error},
+           session_id: session.id
+         }}
+      )
+
+      # The persisted meta["error_message"] is surfaced to the frontend instead
+      # of a generic "Failed to send".
+      assert_broadcast "message_error", %{
+        message_id: ^message_id,
+        status: "error",
+        error: "Token invalid or limit reached. Contact your administrator."
+      }
+    end
+
     test "broadcasts message_error on failed status", %{
       socket: socket,
       job: job,