diff --git a/CHANGELOG.md b/CHANGELOG.md
index 699ff41d0a..499bd0022a 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/assets/js/collaborative-editor/components/MessageList.tsx b/assets/js/collaborative-editor/components/MessageList.tsx
index 48e75a9fca..3cddca7fdd 100644
--- a/assets/js/collaborative-editor/components/MessageList.tsx
+++ b/assets/js/collaborative-editor/components/MessageList.tsx
@@ -652,7 +652,8 @@ export function MessageList({
>
- Failed to send message. Please try again.
+ {message.error ||
+ 'Failed to send message. Please try again.'}
{onRetryMessage && (
- Failed to send
+ {message.error || 'Failed to send'}
{onRetryMessage && (
{
* 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;
}
diff --git a/assets/js/collaborative-editor/types/ai-assistant.ts b/assets/js/collaborative-editor/types/ai-assistant.ts
index 0afb9deae0..1f0405b2bd 100644
--- a/assets/js/collaborative-editor/types/ai-assistant.ts
+++ b/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: (
diff --git a/lib/lightning/ai_assistant/ai_assistant.ex b/lib/lightning/ai_assistant/ai_assistant.ex
index b95886e1a4..5b118ee8fe 100644
--- a/lib/lightning/ai_assistant/ai_assistant.ex
+++ b/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"]
diff --git a/lib/lightning/ai_assistant/message_processor.ex b/lib/lightning/ai_assistant/message_processor.ex
index 735c280d17..f8c196accd 100644
--- a/lib/lightning/ai_assistant/message_processor.ex
+++ b/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
diff --git a/lib/lightning_web/channels/ai_assistant_channel.ex b/lib/lightning_web/channels/ai_assistant_channel.ex
index 1e7e920d16..09fa2a3751 100644
--- a/lib/lightning_web/channels/ai_assistant_channel.ex
+++ b/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
diff --git a/test/lightning/ai_assistant/message_processor_test.exs b/test/lightning/ai_assistant/message_processor_test.exs
index de903cfc88..d6b280f0e2 100644
--- a/test/lightning/ai_assistant/message_processor_test.exs
+++ b/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
diff --git a/test/lightning_web/channels/ai_assistant_channel_test.exs b/test/lightning_web/channels/ai_assistant_channel_test.exs
index a0637d988e..69ac391b62 100644
--- a/test/lightning_web/channels/ai_assistant_channel_test.exs
+++ b/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,