diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java index ed3e8fd80d..d6e3e735eb 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java @@ -16,11 +16,16 @@ package io.agentscope.core.studio; import io.agentscope.core.agent.AgentBase; +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.StreamOptions; +import io.agentscope.core.agent.StreamableAgent; +import io.agentscope.core.message.Msg; import io.agentscope.core.tracing.TracerRegistry; import io.agentscope.core.tracing.telemetry.TelemetryTracer; import java.net.URI; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; /** @@ -121,6 +126,32 @@ public static void shutdown() { wsClient = null; } + /** + * Stream agent events to Studio for real-time visualization and final message persistence. + * + *

This method invokes {@link StreamableAgent#stream(Msg, StreamOptions)} on the provided + * agent to obtain a {@link Flux} of {@link Event} instances, then forwards that stream to + * Studio via a {@link StudioStreamingBridge}. The bridge: + * @param agent the {@link StreamableAgent} to run; must not be {@code null} + * @param input the input {@link Msg} for this invocation (may be {@code null} for agents that + * do not require an initial message) + * @param options optional {@link StreamOptions} controlling the agent's streaming behavior; + * when {@code null}, {@link StreamOptions#defaults()} is used + * @return a {@link Mono} that completes when the event stream has been fully forwarded to + * Studio and, if a final message exists, persisted; if Studio is not initialized, + * this returns a {@link Mono#error(Throwable)} with {@link IllegalStateException} + */ + public static Mono streamToStudio( + StreamableAgent agent, Msg input, StreamOptions options) { + if (!isInitialized() || client == null || wsClient == null) { + return Mono.error(new IllegalStateException("Studio is not initialized")); + } + StreamOptions effective = options != null ? options : StreamOptions.defaults(); + Flux events = agent.stream(input, effective); + StudioStreamingBridge bridge = new StudioStreamingBridge(client, wsClient); + return bridge.forwardToStudio(events); + } + /** * Builder for configuring and initializing Studio integration. */ diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingBridge.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingBridge.java new file mode 100644 index 0000000000..0dd953c304 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingBridge.java @@ -0,0 +1,186 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.studio; + +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.EventType; +import io.agentscope.core.message.Msg; +import java.util.Objects; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +/** + * Bridges agent streaming events to AgentScope Studio. + * + *

This component subscribes to a {@link reactor.core.publisher.Flux} of + * {@link io.agentscope.core.agent.Event} instances produced by a + * {@code StreamableAgent} and forwards them to Studio using both the HTTP + * {@link StudioClient} and the WebSocket-based {@link StudioWebSocketClient}. + * + *

The bridge forwards REASONING, TOOL_RESULT, and SUMMARY events directly to the + * Studio frontend via {@link StudioWebSocketClient#sendStreamEvent(Event)} so that + * intermediate and final reasoning/tool outputs are visible in real time. Non-terminal + * {@link EventType#AGENT_RESULT} events are also streamed as incremental chunks. + * + *

The final {@link EventType#AGENT_RESULT} event with {@code isLast == true} is treated as + * a control signal: its {@link Msg} payload is cached for persistence and it triggers + * {@link StudioWebSocketClient#sendStreamCompleted()} but is not forwarded again as a + * streaming chunk. After the stream completes, the selected final {@link Msg} (terminal + * AGENT_RESULT if present, otherwise the last non-terminal AGENT_RESULT) is persisted once + * via {@link StudioClient#pushMessage(Msg)}. + */ +public class StudioStreamingBridge { + + private static final Logger logger = LoggerFactory.getLogger(StudioStreamingBridge.class); + + private final StudioClient studioClient; + private final StudioWebSocketClient webSocketClient; + + /** + * Creates a new streaming bridge for forwarding agent events to Studio. + * + * @param studioClient HTTP client used to persist the final agent message to Studio + * @param webSocketClient WebSocket client used to stream incremental events and + * completion signals to the Studio frontend + */ + public StudioStreamingBridge(StudioClient studioClient, StudioWebSocketClient webSocketClient) { + this.studioClient = Objects.requireNonNull(studioClient, "studioClient must not be null"); + this.webSocketClient = + Objects.requireNonNull(webSocketClient, "webSocketClient must not be null"); + } + + /** + * Forwards a stream of agent events to Studio for real-time visualization and + * final message persistence. + * + *

    + *
  1. All REASONING, TOOL_RESULT, and SUMMARY events are forwarded to the frontend via + * {@link StudioWebSocketClient#sendStreamEvent(Event)} regardless of the {@code isLast} + * flag so that final reasoning/tool outputs remain visible even when emitted as + * terminal events.
    2. + *
  2. Non-terminal AGENT_RESULT events ({@code isLast == false}) are also forwarded as + * streaming chunks.
    3. + *
  3. The terminal AGENT_RESULT event with {@code isLast == true} is used to: + * + * It is not sent again as a streaming chunk.
    4. + *
  4. If a final {@link Msg} exists, it is persisted once via + * {@link StudioClient#pushMessage(Msg)} after the stream completes.
    5. + *
+ * + *

Note: {@code isLast} is treated purely as a control flag for AGENT_RESULT and does not + * generate additional visible content.

+ */ + public Mono forwardToStudio(Flux eventFlux) { + return Mono.defer( + () -> { + FinalMsgHolder holder = new FinalMsgHolder(); + + return eventFlux + .doOnNext(event -> handleEvent(event, holder)) + .doOnError( + ex -> + logger.error( + "Error occurred during streaming to Studio", + ex)) + .doFinally( + signalType -> { + if (!holder.completedSignalSent) { + try { + webSocketClient.sendStreamCompleted(); + } catch (Exception e) { + logger.error( + "Failed to send streamCompleted in" + + " doFinally", + e); + } + holder.completedSignalSent = true; + } + }) + .then( + Mono.defer( + () -> { + Msg finalMsg = holder.getEffectiveFinalMsg(); + if (finalMsg == null) { + logger.debug( + "No final Msg determined from stream;" + + " skip pushMessage"); + return Mono.empty(); + } + return studioClient.pushMessage(finalMsg); + })); + }); + } + + /** + * Handles individual events from the stream. + */ + private void handleEvent(Event event, FinalMsgHolder holder) { + if (event == null) { + return; + } + + EventType type = event.getType(); + Msg msg = event.getMessage(); + boolean isLast = event.isLast(); + + if (type == EventType.AGENT_RESULT && msg != null) { + holder.lastAgentResultMsg = msg; + } + + if (type == EventType.AGENT_RESULT && isLast) { + if (msg != null) { + holder.terminalAgentResultMsg = msg; + } + if (!holder.completedSignalSent) { + webSocketClient.sendStreamCompleted(); + holder.completedSignalSent = true; + } + return; + } + + if (type == EventType.AGENT_RESULT && !isLast) { + webSocketClient.sendStreamEvent(event); + return; + } + + if (type == EventType.REASONING + || type == EventType.TOOL_RESULT + || type == EventType.SUMMARY) { + webSocketClient.sendStreamEvent(event); + } + } + + private static class FinalMsgHolder { + private Msg terminalAgentResultMsg; + private Msg lastAgentResultMsg; + private boolean completedSignalSent; + + Msg getEffectiveFinalMsg() { + if (terminalAgentResultMsg != null) { + return terminalAgentResultMsg; + } + if (lastAgentResultMsg != null) { + return lastAgentResultMsg; + } + return null; + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioWebSocketClient.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioWebSocketClient.java index 57086442d0..c67c32d1fa 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioWebSocketClient.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioWebSocketClient.java @@ -15,7 +15,13 @@ */ package io.agentscope.core.studio; +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.EventType; import io.agentscope.core.message.ContentBlock; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ThinkingBlock; +import io.agentscope.core.message.ToolResultBlock; import io.agentscope.core.util.JsonUtils; import io.socket.client.IO; import io.socket.client.Socket; @@ -395,4 +401,127 @@ public Map getStructuredInput() { return structuredInput; } } + + /** + * Sends a stream event to Studio in real-time. + * + *

This method delivers best-effort, non-guaranteed notifications over the + * existing Socket.IO connection. If the underlying WebSocket is {@code null} + * or not connected, the method logs a warning and returns without throwing + * an exception and without attempting any reconnection. + * + * @param event the agent event to forward to Studio; if {@code null}, the + * call is ignored and no payload is emitted + */ + public void sendStreamEvent(Event event) { + if (socket == null || !socket.connected()) { + logger.warn("Attempted to send stream event while WebSocket is not connected"); + return; + } + if (event == null) { + return; + } + + try { + JSONObject payload = new JSONObject(); + EventType type = event.getType(); + payload.put("eventType", type != null ? type.name() : "UNKNOWN"); + payload.put("isLast", event.isLast()); + + Msg msg = event.getMessage(); + if (msg != null) { + String text = extractTextFromMsg(msg); + if (text != null) { + payload.put("text", text); + } + if (msg.getRole() != null) { + payload.put("role", msg.getRole().name()); + } + if (msg.getName() != null) { + payload.put("name", msg.getName()); + } + if (msg.getMetadata() != null) { + String metadataJson = JsonUtils.getJsonCodec().toJson(msg.getMetadata()); + payload.put("metadata", new JSONObject(metadataJson)); + } + } + + socket.emit("studioStreamEvent", payload); + } catch (Exception e) { + logger.error("Failed to send stream event to Studio", e); + } + } + + /** + * Sends a stream completed event to Studio. + * + *

This method delivers a best-effort completion signal over the + * existing Socket.IO connection. If the underlying WebSocket is {@code null} + * or not connected, the method logs a warning and returns without throwing + * an exception and without attempting any reconnection. + */ + public void sendStreamCompleted() { + if (socket == null || !socket.connected()) { + logger.warn("Attempted to send streamCompleted while WebSocket is not connected"); + return; + } + try { + JSONObject payload = new JSONObject(); + payload.put("status", "completed"); + socket.emit("studioStreamCompleted", payload); + } catch (Exception e) { + logger.error("Failed to send streamCompleted event to Studio", e); + } + } + + /** + * Extracts text content from a Msg for streaming to Studio. + */ + private String extractTextFromMsg(Msg msg) { + if (msg == null || msg.getContent() == null || msg.getContent().isEmpty()) { + return null; + } + + StringBuilder builder = new StringBuilder(); + for (Object block : msg.getContent()) { + if (!(block instanceof ContentBlock)) { + continue; + } + if (block instanceof TextBlock) { + TextBlock tb = (TextBlock) block; + if (tb.getText() != null && !tb.getText().isEmpty()) { + appendWithNewline(builder, tb.getText()); + } + } else if (block instanceof ThinkingBlock) { + ThinkingBlock tb = (ThinkingBlock) block; + if (tb.getThinking() != null && !tb.getThinking().isEmpty()) { + appendWithNewline(builder, tb.getThinking()); + } + } else if (block instanceof ToolResultBlock) { + ToolResultBlock trb = (ToolResultBlock) block; + for (ContentBlock outputBlock : trb.getOutput()) { + if (outputBlock instanceof TextBlock) { + TextBlock otb = (TextBlock) outputBlock; + if (otb.getText() != null && !otb.getText().isEmpty()) { + appendWithNewline(builder, otb.getText()); + } + } else if (outputBlock instanceof ThinkingBlock) { + ThinkingBlock otb = (ThinkingBlock) outputBlock; + if (otb.getThinking() != null && !otb.getThinking().isEmpty()) { + appendWithNewline(builder, otb.getThinking()); + } + } + } + } + } + + return builder.length() > 0 ? builder.toString() : null; + } + + private void appendWithNewline(StringBuilder builder, String text) { + if (builder.length() > 0) { + builder.append('\n'); + } + builder.append(text); + } } diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioManagerTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioManagerTest.java index 51156bcbd1..20b510f810 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioManagerTest.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioManagerTest.java @@ -15,13 +15,31 @@ */ package io.agentscope.core.studio; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; - +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.EventType; +import io.agentscope.core.agent.StreamOptions; +import io.agentscope.core.agent.StreamableAgent; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import java.lang.reflect.Field; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.test.StepVerifier; @DisplayName("StudioManager Tests") class StudioManagerTest { @@ -125,4 +143,89 @@ void testShutdownIdempotent() { assertNull(StudioManager.getClient()); assertFalse(StudioManager.isInitialized()); } + + @Test + @DisplayName("streamToStudio() should return error when Studio is not initialized") + void streamToStudio_WhenNotInitialized_ShouldReturnError() { + StreamableAgent mockAgent = mock(StreamableAgent.class); + Msg input = + Msg.builder() + .role(MsgRole.USER) + .content(TextBlock.builder().text("test").build()) + .build(); + + StepVerifier.create(StudioManager.streamToStudio(mockAgent, input, null)) + .expectErrorSatisfies( + throwable -> { + assertTrue(throwable instanceof IllegalStateException); + assertEquals("Studio is not initialized", throwable.getMessage()); + }) + .verify(); + } + + @Test + @DisplayName("streamToStudio() should use defaults when options is null") + void streamToStudio_WithNullOptions_ShouldUseDefaults() throws Exception { + // Setup: Use reflection to set static fields to simulate initialized state + StudioClient mockClient = mock(StudioClient.class); + StudioWebSocketClient mockWsClient = mock(StudioWebSocketClient.class); + + setStaticField("client", mockClient); + setStaticField("wsClient", mockWsClient); + setStaticField("config", StudioConfig.builder().build()); + + StreamableAgent mockAgent = mock(StreamableAgent.class); + Msg input = + Msg.builder() + .role(MsgRole.USER) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, input, true); + + when(mockAgent.stream(any(Msg.class), any(StreamOptions.class))) + .thenReturn(Flux.just(event)); + when(mockClient.pushMessage(any(Msg.class))).thenReturn(Flux.empty().then()); + + StepVerifier.create(StudioManager.streamToStudio(mockAgent, input, null)).verifyComplete(); + + // Verify that stream was called with some StreamOptions (defaults) + verify(mockAgent, times(1)).stream(any(Msg.class), any(StreamOptions.class)); + } + + @Test + @DisplayName("streamToStudio() should use provided options when not null") + void streamToStudio_WithCustomOptions_ShouldUseCustomOptions() throws Exception { + // Setup: Use reflection to set static fields to simulate initialized state + StudioClient mockClient = mock(StudioClient.class); + StudioWebSocketClient mockWsClient = mock(StudioWebSocketClient.class); + + setStaticField("client", mockClient); + setStaticField("wsClient", mockWsClient); + setStaticField("config", StudioConfig.builder().build()); + + StreamableAgent mockAgent = mock(StreamableAgent.class); + Msg input = + Msg.builder() + .role(MsgRole.USER) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, input, true); + StreamOptions customOptions = + StreamOptions.builder().eventTypes(EventType.REASONING).build(); + + when(mockAgent.stream(eq(input), eq(customOptions))).thenReturn(Flux.just(event)); + when(mockClient.pushMessage(any(Msg.class))).thenReturn(Flux.empty().then()); + + StepVerifier.create(StudioManager.streamToStudio(mockAgent, input, customOptions)) + .verifyComplete(); + + // Verify that stream was called with the custom options + verify(mockAgent, times(1)).stream(eq(input), eq(customOptions)); + } + + private void setStaticField(String fieldName, Object value) throws Exception { + Field field = StudioManager.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(null, value); + } } diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java new file mode 100644 index 0000000000..fe3dbe80d9 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java @@ -0,0 +1,165 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.studio; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.argThat; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.EventType; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ThinkingBlock; +import io.agentscope.core.message.ToolResultBlock; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +@ExtendWith(MockitoExtension.class) +class StudioStreamingBridgeTest { + + private Msg msg(String text) { + return Msg.builder() + .role(MsgRole.ASSISTANT) + .name("Agent") + .content(TextBlock.builder().text(text).build()) + .build(); + } + + private Msg thinkingMsg(String thinking) { + return Msg.builder() + .role(MsgRole.ASSISTANT) + .name("Agent") + .content(ThinkingBlock.builder().thinking(thinking).build()) + .build(); + } + + private Msg toolResultMsg(String text) { + ToolResultBlock toolResult = ToolResultBlock.text(text); + return Msg.builder().role(MsgRole.ASSISTANT).name("Agent").content(toolResult).build(); + } + + @Test + void testStreamingAndCompletion() { + + StudioClient mockClient = Mockito.mock(StudioClient.class); + StudioWebSocketClient mockWs = Mockito.mock(StudioWebSocketClient.class); + + when(mockClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + StudioStreamingBridge bridge = new StudioStreamingBridge(mockClient, mockWs); + + Event reasoning1 = new Event(EventType.REASONING, msg("thinking 1"), false); + Event reasoning2 = new Event(EventType.REASONING, msg("thinking 2"), false); + + Event tool1 = new Event(EventType.TOOL_RESULT, msg("tool out 1"), false); + + Event answerChunk1 = new Event(EventType.AGENT_RESULT, msg("hello "), false); + Event answerChunk2 = new Event(EventType.AGENT_RESULT, msg("world"), false); + + Msg finalMsg = msg("hello world"); + Event answerFinal = new Event(EventType.AGENT_RESULT, finalMsg, true); + + Flux flux = + Flux.fromIterable( + List.of( + reasoning1, + reasoning2, + tool1, + answerChunk1, + answerChunk2, + answerFinal)); + + bridge.forwardToStudio(flux).block(); + + verify(mockWs, times(5)).sendStreamEvent(any(Event.class)); + + verify(mockWs, Mockito.atLeastOnce()) + .sendStreamEvent( + argThat( + e -> + !e.isLast() + && (e.getType() == EventType.REASONING + || e.getType() == EventType.TOOL_RESULT + || e.getType() == EventType.AGENT_RESULT))); + + verify(mockWs, times(1)).sendStreamCompleted(); + + verify(mockClient, times(1)).pushMessage(eq(finalMsg)); + } + + @Test + void testFallbackWhenNoTerminalAgentResult() { + StudioClient mockClient = Mockito.mock(StudioClient.class); + StudioWebSocketClient mockWs = Mockito.mock(StudioWebSocketClient.class); + + when(mockClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + StudioStreamingBridge bridge = new StudioStreamingBridge(mockClient, mockWs); + + Event chunk1 = new Event(EventType.AGENT_RESULT, msg("foo"), false); + Msg lastMsg = msg("foo bar"); + Event chunk2 = new Event(EventType.AGENT_RESULT, lastMsg, false); + + Flux flux = Flux.just(chunk1, chunk2); + + bridge.forwardToStudio(flux).block(); + + verify(mockWs, times(2)).sendStreamEvent(any(Event.class)); + verify(mockWs, times(1)).sendStreamCompleted(); + verify(mockClient, times(1)).pushMessage(eq(lastMsg)); + } + + @Test + void testReasoningAndToolResultEventsWithThinkingAndToolResultBlocksAreStreamed() { + StudioClient mockClient = Mockito.mock(StudioClient.class); + StudioWebSocketClient mockWs = Mockito.mock(StudioWebSocketClient.class); + + when(mockClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + StudioStreamingBridge bridge = new StudioStreamingBridge(mockClient, mockWs); + + Event reasoningFinal = + new Event(EventType.REASONING, thinkingMsg("chain-of-thought"), true); + Event toolFinal = new Event(EventType.TOOL_RESULT, toolResultMsg("tool-output"), true); + + // Include a terminal AGENT_RESULT so that the stream completes normally + Msg finalMsg = msg("answer"); + Event answerFinal = new Event(EventType.AGENT_RESULT, finalMsg, true); + + Flux flux = Flux.just(reasoningFinal, toolFinal, answerFinal); + + bridge.forwardToStudio(flux).block(); + + // Both reasoning and tool result events (even with isLast == true) should be streamed + verify(mockWs, times(1)).sendStreamEvent(reasoningFinal); + verify(mockWs, times(1)).sendStreamEvent(toolFinal); + // Terminal AGENT_RESULT should not be streamed as an event, only used for completion + verify(mockWs, Mockito.never()).sendStreamEvent(answerFinal); + + verify(mockWs, times(1)).sendStreamCompleted(); + verify(mockClient, times(1)).pushMessage(eq(finalMsg)); + } +} diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioWebSocketClientTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioWebSocketClientTest.java index 596b59df3a..a196ae2bd0 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioWebSocketClientTest.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioWebSocketClientTest.java @@ -18,19 +18,34 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; - +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.Event; +import io.agentscope.core.agent.EventType; import io.agentscope.core.message.ContentBlock; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ThinkingBlock; +import io.agentscope.core.message.ToolResultBlock; import io.socket.client.Socket; import java.time.Duration; import java.util.List; import java.util.Map; import org.json.JSONArray; +import org.json.JSONException; import org.json.JSONObject; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -304,4 +319,312 @@ void testConstructorWithSocket() { assertEquals(true, clientWithSocket.isConnected()); clientWithSocket.close(); } + + @Test + @DisplayName("sendStreamEvent should log warning when socket is null") + void sendStreamEvent_WithNullSocket_ShouldLogWarning() { + // Client created without socket - should not throw + StudioWebSocketClient newClient = new StudioWebSocketClient(config); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + newClient.sendStreamEvent(event); + } + + @Test + @DisplayName("sendStreamEvent should log warning when socket is not connected") + void sendStreamEvent_WithDisconnectedSocket_ShouldLogWarning() { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(false); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + verify(mockSocket, never()).emit(any(), any()); + } + + @Test + @DisplayName("sendStreamEvent should return when event is null") + void sendStreamEvent_WithNullEvent_ShouldReturn() { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + + // Should not throw + socketClient.sendStreamEvent(null); + + verify(mockSocket, never()).emit(any(), any()); + } + + @Test + @DisplayName("sendStreamEvent should emit event with TextBlock content") + void sendStreamEvent_WithTextBlock_ShouldEmit() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .name("TestAgent") + .content(TextBlock.builder().text("Hello World").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, true); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + JSONObject payload = payloadCaptor.getValue(); + assertEquals("AGENT_RESULT", payload.getString("eventType")); + assertTrue(payload.getBoolean("isLast")); + assertEquals("Hello World", payload.getString("text")); + assertEquals("ASSISTANT", payload.getString("role")); + assertEquals("TestAgent", payload.getString("name")); + } + + @Test + @DisplayName("sendStreamEvent should emit event with ThinkingBlock content") + void sendStreamEvent_WithThinkingBlock_ShouldEmit() throws JSONException { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(ThinkingBlock.builder().thinking("Thinking process").build()) + .build(); + Event event = new Event(EventType.REASONING, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + JSONObject payload = payloadCaptor.getValue(); + assertEquals("REASONING", payload.getString("eventType")); + assertEquals("Thinking process", payload.getString("text")); + } + + @Test + @DisplayName("sendStreamEvent should emit event with ToolResultBlock content") + void sendStreamEvent_WithToolResultBlock_ShouldEmit() throws JSONException { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + ToolResultBlock toolResult = ToolResultBlock.text("Tool output"); + Msg msg = Msg.builder().role(MsgRole.ASSISTANT).content(toolResult).build(); + Event event = new Event(EventType.TOOL_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + JSONObject payload = payloadCaptor.getValue(); + assertEquals("TOOL_RESULT", payload.getString("eventType")); + assertEquals("Tool output", payload.getString("text")); + } + + @Test + @DisplayName("sendStreamEvent should handle event with null type") + void sendStreamEvent_WithNullType_ShouldUseUnknown() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(null, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + assertEquals("UNKNOWN", payloadCaptor.getValue().getString("eventType")); + } + + @Test + @DisplayName("sendStreamEvent should handle event with null msg") + void sendStreamEvent_WithNullMsg_ShouldEmitWithoutMsgFields() throws JSONException { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Event event = new Event(EventType.AGENT_RESULT, null, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + JSONObject payload = payloadCaptor.getValue(); + assertEquals("AGENT_RESULT", payload.getString("eventType")); + assertFalse(payload.has("text")); + assertFalse(payload.has("role")); + } + + @Test + @DisplayName("sendStreamEvent should handle msg with null name") + void sendStreamEvent_WithNullName_ShouldSkipName() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("test").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + assertFalse(payloadCaptor.getValue().has("name")); + } + + @Test + @DisplayName("sendStreamEvent should handle TextBlock with null text") + void sendStreamEvent_WithNullText_ShouldSkipText() { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text(null).build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + assertFalse(payloadCaptor.getValue().has("text")); + } + + @Test + @DisplayName("sendStreamEvent should handle TextBlock with empty text") + void sendStreamEvent_WithEmptyText_ShouldSkipText() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + assertFalse(payloadCaptor.getValue().has("text")); + } + + @Test + @DisplayName("sendStreamEvent should handle multiple content blocks") + void sendStreamEvent_WithMultipleBlocks_ShouldConcatenate() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + Msg msg = + Msg.builder() + .role(MsgRole.ASSISTANT) + .content( + TextBlock.builder().text("First").build(), + TextBlock.builder().text("Second").build(), + ThinkingBlock.builder().thinking("Thinking").build()) + .build(); + Event event = new Event(EventType.AGENT_RESULT, msg, false); + + socketClient.sendStreamEvent(event); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + String text = payloadCaptor.getValue().getString("text"); + assertTrue(text.contains("First")); + assertTrue(text.contains("Second")); + assertTrue(text.contains("Thinking")); + } + + @Test + @DisplayName("sendStreamCompleted should log warning when socket is null") + void sendStreamCompleted_WithNullSocket_ShouldLogWarning() { + StudioWebSocketClient newClient = new StudioWebSocketClient(config); + + // Should not throw + newClient.sendStreamCompleted(); + } + + @Test + @DisplayName("sendStreamCompleted should log warning when socket is not connected") + void sendStreamCompleted_WithDisconnectedSocket_ShouldLogWarning() { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(false); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + + // Should not throw + socketClient.sendStreamCompleted(); + + verify(mockSocket, never()).emit(any(), any()); + } + + @Test + @DisplayName("sendStreamCompleted should emit completed event") + void sendStreamCompleted_WithConnectedSocket_ShouldEmit() throws Exception { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + + socketClient.sendStreamCompleted(); + + ArgumentCaptor payloadCaptor = ArgumentCaptor.forClass(JSONObject.class); + verify(mockSocket, times(1)).emit(any(), payloadCaptor.capture()); + + assertEquals("completed", payloadCaptor.getValue().getString("status")); + } + + @Test + @DisplayName("sendStreamCompleted should handle emit exception gracefully") + void sendStreamCompleted_WithEmitException_ShouldLogError() { + Socket mockSocket = mock(Socket.class); + when(mockSocket.connected()).thenReturn(true); + when(mockSocket.emit(any(), any())).thenThrow(new RuntimeException("Emit failed")); + + StudioWebSocketClient socketClient = new StudioWebSocketClient(config, mockSocket); + + // Should not throw + socketClient.sendStreamCompleted(); + } }