Skip to content

Handle AI provider errors in WebSocket chat and clear loading state #41

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rdehuyss merged 2 commits into from
Mar 30, 2026
+169 −25
Merged

Handle AI provider errors in WebSocket chat and clear loading state #41

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c29a0c9
Handle AI provider errors in WebSocket chat and clear loading state
arefbehboudi Mar 27, 2026
5d59b43
Small improvements.
rdehuyss Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions app/src/main/java/ai/javaclaw/chat/ChatChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,10 +72,10 @@ public void clearWsSession(WebSocketSession session) {
* Sends a raw HTML fragment to the active WebSocket session.
* Used by the WebSocket handler to push user/agent bubbles and typing indicators.
*/
public void sendHtml(String html) throws IOException {
public void sendHtml(String... html) throws IOException {
WebSocketSession session = wsSession.get();
if (session != null && session.isOpen()) {
session.sendMessage(new TextMessage(html));
session.sendMessage(new TextMessage(String.join(System.lineSeparator(), html)));
}
}

Expand Down
60 changes: 37 additions & 23 deletions app/src/main/java/ai/javaclaw/chat/ws/ChatWebSocketHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,16 +36,15 @@ public void afterConnectionEstablished(WebSocketSession session) throws Exceptio
log.info("WebChat WebSocket connected: {}", session.getId());

List<String> ids = chatChannel.conversationIds();
String selectedId = ids.get(0);
String selectedId = ids.getFirst();

String selector = ChatHtml.conversationSelector(ids, selectedId);
String bubbles = String.join("", chatChannel.loadHistoryAsHtml(selectedId));
String conversationSelector = ChatHtml.conversationSelector(ids, selectedId);
String bubbles = String.join(System.lineSeparator(), chatChannel.loadHistoryAsHtml(selectedId));
String inputArea = ChatHtml.chatInputArea(selectedId);
session.sendMessage(new TextMessage(
Htmx.oobInnerHtml("channel-selector", selector) +
Htmx.oobInnerHtml("chat-messages", bubbles) +
Htmx.oobInnerHtml("chat-input-area", inputArea)
));
chatChannel.sendHtml(
Htmx.oobInnerHtml("channel-selector", conversationSelector),
Htmx.oobInnerHtml("chat-messages", bubbles),
Htmx.oobInnerHtml("chat-input-area", inputArea));
}

@Override
Expand All @@ -71,12 +70,11 @@ private void handleChannelChanged(Map<String, Object> payload) throws Exception
String conversationId = (String) payload.get("conversationId");
if (conversationId == null || conversationId.isBlank()) return;

String bubbles = String.join("", chatChannel.loadHistoryAsHtml(conversationId));
String bubbles = String.join(System.lineSeparator(), chatChannel.loadHistoryAsHtml(conversationId));
String inputArea = ChatHtml.chatInputArea(conversationId);
chatChannel.sendHtml(
Htmx.oobInnerHtml("chat-messages", bubbles) +
Htmx.oobInnerHtml("chat-input-area", inputArea)
);
Htmx.oobInnerHtml("chat-messages", bubbles),
Htmx.oobInnerHtml("chat-input-area", inputArea));
}

private void handleUserMessage(Map<String, Object> payload) throws Exception {
Expand All @@ -89,17 +87,33 @@ private void handleUserMessage(Map<String, Object> payload) throws Exception {

// Echo user message + show typing indicator
chatChannel.sendHtml(
Htmx.oobAppend("chat-messages", ChatHtml.userBubble(userMessage)) +
Htmx.oobReplace("typing-indicator", ChatHtml.typingDots())
);
Htmx.oobAppend("chat-messages", ChatHtml.userBubble(userMessage)),
Htmx.oobReplace("typing-indicator", ChatHtml.typingDots()));

try {
// Call agent (blocking — background tasks may push messages via ChatChannel during this)
String response = chatChannel.chat(conversationId, userMessage);
chatChannel.sendHtml(
Htmx.oobAppend("chat-messages", ChatHtml.agentBubble(response)),
Htmx.oobReplace("typing-indicator", ""));
} catch (RuntimeException ex) {
log.warn("Chat request failed for conversation {}", conversationId, ex);
chatChannel.sendHtml(
Htmx.oobAppend("chat-messages", ChatHtml.agentBubble(genericUserFacingError(ex))),
Htmx.oobReplace("typing-indicator", ""));
}
}

// Call agent (blocking — background tasks may push messages via ChatChannel during this)
String response = chatChannel.chat(conversationId, userMessage);
private static String genericUserFacingError(RuntimeException ex) {
return "An error occurred while contacting the AI provider.\nDetails: " + summarizeError(ex);
}

// Send agent response + clear typing indicator
chatChannel.sendHtml(
Htmx.oobAppend("chat-messages", ChatHtml.agentBubble(response)) +
Htmx.oobReplace("typing-indicator", "")
);
private static String summarizeError(Throwable ex) {
String message = ex.getMessage();
if (message == null || message.isBlank()) {
return ex.getClass().getSimpleName();
}

return message;
}
}
}
130 changes: 130 additions & 0 deletions app/src/test/java/ai/javaclaw/chat/ws/ChatWebSocketHandlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
package ai.javaclaw.chat.ws;

import ai.javaclaw.chat.ChatChannel;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import tools.jackson.databind.ObjectMapper;

import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

class ChatWebSocketHandlerTest {

@Test
void handleUserMessageShowsNonTransientErrorAndClearsTypingIndicatorWhenAgentFails() throws Exception {
ChatChannel chatChannel = mock(ChatChannel.class);
WebSocketSession session = mock(WebSocketSession.class);
ChatWebSocketHandler handler = new ChatWebSocketHandler(chatChannel, new ObjectMapper());

when(chatChannel.chat("web", "hello")).thenThrow(new RuntimeException("""
HTTP 401 - {
"error": {
"message": "Incorrect API key provided: Test.",
"code": "invalid_api_key"
}
}
"""));

handler.handleTextMessage(session, new TextMessage(new ObjectMapper().writeValueAsString(Map.of(
"type", "userMessage",
"conversationId", "web",
"message", "hello"
))));

ArgumentCaptor<String[]> htmlCaptor = ArgumentCaptor.forClass(String[].class);
var inOrder = inOrder(chatChannel);
inOrder.verify(chatChannel).sendHtml(htmlCaptor.capture());
inOrder.verify(chatChannel).chat("web", "hello");
inOrder.verify(chatChannel).sendHtml(htmlCaptor.capture());
verifyNoMoreInteractions(chatChannel);

assertThat(String.join("", htmlCaptor.getAllValues().get(0)))
.contains("hello")
.contains("typing-indicator")
.contains("ar-typing");

assertThat(String.join("", htmlCaptor.getAllValues().get(1)))
.contains("An error occurred while contacting the AI provider")
.contains("Details: HTTP 401 - {")
.contains("typing-indicator")
.doesNotContain("ar-typing");
}

@Test
void handleUserMessageShowsGenericProviderErrorForUnexpectedFailures() throws Exception {
ChatChannel chatChannel = mock(ChatChannel.class);
WebSocketSession session = mock(WebSocketSession.class);
ChatWebSocketHandler handler = new ChatWebSocketHandler(chatChannel, new ObjectMapper());

when(chatChannel.chat(anyString(), anyString())).thenThrow(new RuntimeException("boom"));

handler.handleTextMessage(session, new TextMessage(new ObjectMapper().writeValueAsString(Map.of(
"type", "userMessage",
"conversationId", "web",
"message", "hello"
))));

ArgumentCaptor<String[]> htmlCaptor = ArgumentCaptor.forClass(String[].class);
var inOrder = inOrder(chatChannel);
inOrder.verify(chatChannel).sendHtml(htmlCaptor.capture());
inOrder.verify(chatChannel).chat("web", "hello");
inOrder.verify(chatChannel).sendHtml(htmlCaptor.capture());

assertThat(String.join("", htmlCaptor.getAllValues().get(1)))
.contains("An error occurred while contacting the AI provider")
.contains("Details: boom");
}

@Test
void handleChannelChangedSendsHistoryAndInputArea() throws Exception {
ChatChannel chatChannel = mock(ChatChannel.class);
WebSocketSession session = mock(WebSocketSession.class);
ChatWebSocketHandler handler = new ChatWebSocketHandler(chatChannel, new ObjectMapper());

when(chatChannel.loadHistoryAsHtml("web")).thenReturn(List.of("<div>history</div>"));

handler.handleTextMessage(session, new TextMessage(new ObjectMapper().writeValueAsString(Map.of(
"type", "channelChanged",
"conversationId", "web"
))));

ArgumentCaptor<String[]> htmlCaptor = ArgumentCaptor.forClass(String[].class);
verify(chatChannel).sendHtml(htmlCaptor.capture());

assertThat(String.join("", htmlCaptor.getValue()))
.contains("chat-messages")
.contains("history")
.contains("chat-input-area");
}

@Test
void afterConnectionEstablishedSendsSelectorHistoryAndInputArea() throws Exception {
ChatChannel chatChannel = mock(ChatChannel.class);
WebSocketSession session = mock(WebSocketSession.class);
ChatWebSocketHandler handler = new ChatWebSocketHandler(chatChannel, new ObjectMapper());

when(chatChannel.conversationIds()).thenReturn(List.of("web"));
when(chatChannel.loadHistoryAsHtml("web")).thenReturn(List.of("<div>history</div>"));

handler.afterConnectionEstablished(session);

ArgumentCaptor<String[]> htmlCaptor = ArgumentCaptor.forClass(String[].class);
verify(chatChannel).sendHtml(htmlCaptor.capture());

assertThat(String.join("", htmlCaptor.getValue()))
.contains("channel-selector")
.contains("chat-messages")
.contains("history")
.contains("chat-input-area");
}
}