From 0cd98ec8fa6f72d1107b3e9f1d557c6d10608f94 Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 15 Jan 2026 00:13:10 +0800 Subject: [PATCH 1/7] feat: add streaming output for studio --- .../core/studio/StudioMessageHook.java | 113 +++++++++++++++ .../core/studio/StudioMessageHookTest.java | 129 ++++++++++++++++++ 2 files changed, 242 insertions(+) diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java index b739c2ed57..84961c5eb8 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java @@ -18,11 +18,21 @@ import io.agentscope.core.hook.Hook; import io.agentscope.core.hook.HookEvent; import io.agentscope.core.hook.PostCallEvent; +import io.agentscope.core.hook.ReasoningChunkEvent; +import io.agentscope.core.hook.PostReasoningEvent; +import io.agentscope.core.hook.ActingChunkEvent; +import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.ToolResultBlock; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + /** * Hook for automatically forwarding agent messages to Studio. * @@ -91,6 +101,109 @@ public Mono onEvent(T event) { return Mono.just(event); }); } + + // Reasoning incremental chunk + if (event instanceof ReasoningChunkEvent) { + ReasoningChunkEvent e = (ReasoningChunkEvent) event; + Msg chunk = e.getIncrementalChunk(); + if (chunk != null && studioClient != null) { + Msg tagged = withEventMetadata(chunk, "reasoning", false); + return studioClient + .pushMessage(tagged) + .thenReturn(event) + .onErrorResume( + ex -> { + logger.error("Failed to push reasoning chunk to Studio", ex); + return Mono.just(event); + }); + } + return Mono.just(event); + } + + // Reasoning final result (after stream completes) + if (event instanceof PostReasoningEvent) { + PostReasoningEvent e = (PostReasoningEvent) event; + Msg finalMsg = e.getReasoningMessage(); + if (finalMsg != null && studioClient != null) { + Msg tagged = withEventMetadata(finalMsg, "reasoning", true); + return studioClient + .pushMessage(tagged) + .thenReturn(event) + .onErrorResume( + ex -> { + logger.error("Failed to push final reasoning to Studio", ex); + return Mono.just(event); + }); + } + return Mono.just(event); + } + + // Acting (tool) incremental chunk + if (event instanceof ActingChunkEvent) { + ActingChunkEvent e = (ActingChunkEvent) event; + ToolResultBlock chunk = e.getChunk(); + if (chunk != null && studioClient != null) { + // build a Msg with TOOL role similar to streaming hook + Msg toolMsg = Msg.builder() + .name("tool") + .role(MsgRole.TOOL) + .content(List.of(chunk)) + .build(); + Msg tagged = withEventMetadata(toolMsg, "tool_result", false); + return studioClient + .pushMessage(tagged) + .thenReturn(event) + .onErrorResume( + ex -> { + logger.error("Failed to push acting chunk to Studio", ex); + return Mono.just(event); + }); + } + return Mono.just(event); + } + + // Acting (tool) final result + if (event instanceof PostActingEvent) { + PostActingEvent e = (PostActingEvent) event; + ToolResultBlock result = e.getToolResult(); + if (result != null && studioClient != null) { + Msg toolMsg = Msg.builder() + .name("tool") + .role(MsgRole.TOOL) + .content(List.of(result)) + .build(); + Msg tagged = withEventMetadata(toolMsg, "tool_result", true); + return studioClient + .pushMessage(tagged) + .thenReturn(event) + .onErrorResume( + ex -> { + logger.error("Failed to push tool result to Studio", ex); + return Mono.just(event); + }); + } + return Mono.just(event); + } + return Mono.just(event); } + + /** + * Wrap/augment a Msg with metadata that indicates event type and whether it's the last chunk. + */ + private Msg withEventMetadata(Msg orig, String eventType, boolean isLast) { + Map meta = new HashMap<>(); + if (orig.getMetadata() != null) { + meta.putAll(orig.getMetadata()); + } + meta.put("studio_event_type", eventType); + meta.put("studio_is_last", isLast); + // Rebuild msg copying fields and injecting metadata; adapt if Msg.Builder has different methods + return Msg.builder() + .name(orig.getName()) + .role(orig.getRole()) + .content(orig.getContent()) + .metadata(meta) + .build(); + } } diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java index 48b07bc8b8..fa4a2361ef 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java @@ -17,6 +17,7 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -27,12 +28,21 @@ import io.agentscope.core.hook.HookEvent; import io.agentscope.core.hook.PostCallEvent; import io.agentscope.core.hook.PreCallEvent; +import io.agentscope.core.hook.ReasoningChunkEvent; +import io.agentscope.core.hook.PostReasoningEvent; +import io.agentscope.core.hook.ActingChunkEvent; +import io.agentscope.core.hook.PostActingEvent; +import io.agentscope.core.message.ToolResultBlock; import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.Map; + +import io.agentscope.core.message.ToolUseBlock; +import io.agentscope.core.tool.Toolkit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -229,4 +239,123 @@ void testMultipleEvents() { verify(mockStudioClient, times(1)).pushMessage(msg1); verify(mockStudioClient, times(1)).pushMessage(msg2); } + + @Test + @DisplayName("Should forward reasoning chunk to Studio") + void testReasoningChunkEventForwards() { + // Mock push success + when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + // Build incremental and accumulated msg + Msg incremental = + Msg.builder() + .name("Assistant") + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("partial chunk").build()) + .build(); + Msg accumulated = + Msg.builder() + .name("Assistant") + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("accumulated so far").build()) + .build(); + + ReasoningChunkEvent event = + new ReasoningChunkEvent(mockAgent, "test-model", null, incremental, accumulated); + + Mono result = hook.onEvent(event); + + StepVerifier.create(result).expectNext(event).verifyComplete(); + + verify(mockStudioClient, times(1)) + .pushMessage(argThat(m -> { + Map md = m.getMetadata(); + return md != null + && "reasoning".equals(md.get("studio_event_type")) + && Boolean.FALSE.equals(md.get("studio_is_last")); + })); + } + + @Test + @DisplayName("Should forward final reasoning result to Studio") + void testPostReasoningEventForwards() { + // Mock push success + when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + Msg finalMsg = + Msg.builder() + .name("Assistant") + .role(MsgRole.ASSISTANT) + .content(TextBlock.builder().text("final reasoning").build()) + .build(); + + PostReasoningEvent event = new PostReasoningEvent(mockAgent, "test-model", null, finalMsg); + + Mono result = hook.onEvent(event); + + StepVerifier.create(result).expectNext(event).verifyComplete(); + + verify(mockStudioClient, times(1)) + .pushMessage(argThat(m -> { + Map md = m.getMetadata(); + return md != null + && "reasoning".equals(md.get("studio_event_type")) + && Boolean.TRUE.equals(md.get("studio_is_last")); + })); + } + + @Test + @DisplayName("Should forward acting (tool) chunk to Studio") + void testActingChunkEventForwards() { + // Mock push success + when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + Toolkit toolkit = new Toolkit(); + ToolUseBlock toolUse = ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); + ToolResultBlock chunk = ToolResultBlock.text("progress update"); + + ActingChunkEvent event = new ActingChunkEvent(mockAgent, toolkit, toolUse, chunk); + + Mono result = hook.onEvent(event); + + StepVerifier.create(result).expectNext(event).verifyComplete(); + + verify(mockStudioClient, times(1)) + .pushMessage(argThat(m -> { + // Expect a TOOL role message with metadata and a ToolResultBlock in content + Map md = m.getMetadata(); + boolean hasMeta = md != null && "tool_result".equals(md.get("studio_event_type")) && Boolean.FALSE.equals(md.get("studio_is_last")); + boolean isToolRole = m.getRole() == MsgRole.TOOL; + boolean hasToolResult = + !m.getContentBlocks(io.agentscope.core.message.ToolResultBlock.class).isEmpty(); + return hasMeta && isToolRole && hasToolResult; + })); + } + + @Test + @DisplayName("Should forward final acting (tool) result to Studio") + void testPostActingEventForwards() { + // Mock push success + when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); + + Toolkit toolkit = new Toolkit(); + ToolUseBlock toolUse = ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); + ToolResultBlock resultBlock = ToolResultBlock.text("final result"); + + PostActingEvent event = new PostActingEvent(mockAgent, toolkit, toolUse, resultBlock); + + Mono result = hook.onEvent(event); + + StepVerifier.create(result).expectNext(event).verifyComplete(); + + verify(mockStudioClient, times(1)) + .pushMessage(argThat(m -> { + Map md = m.getMetadata(); + boolean hasMeta = md != null && "tool_result".equals(md.get("studio_event_type")) && Boolean.TRUE.equals(md.get("studio_is_last")); + boolean isToolRole = m.getRole() == MsgRole.TOOL; + boolean hasToolResult = + !m.getContentBlocks(io.agentscope.core.message.ToolResultBlock.class).isEmpty(); + return hasMeta && isToolRole && hasToolResult; + })); + } } From 1e8ac607353a5f19d4951dfa500297d82c2fc04f Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 15 Jan 2026 00:27:10 +0800 Subject: [PATCH 2/7] fix the code format --- .../core/studio/StudioMessageHook.java | 38 +++---- .../core/studio/StudioMessageHookTest.java | 102 +++++++++++------- 2 files changed, 84 insertions(+), 56 deletions(-) diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java index 84961c5eb8..57c25e9f01 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java @@ -15,23 +15,22 @@ */ package io.agentscope.core.studio; +import io.agentscope.core.hook.ActingChunkEvent; import io.agentscope.core.hook.Hook; import io.agentscope.core.hook.HookEvent; +import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.hook.PostCallEvent; -import io.agentscope.core.hook.ReasoningChunkEvent; import io.agentscope.core.hook.PostReasoningEvent; -import io.agentscope.core.hook.ActingChunkEvent; -import io.agentscope.core.hook.PostActingEvent; +import io.agentscope.core.hook.ReasoningChunkEvent; import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.ToolResultBlock; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import reactor.core.publisher.Mono; - import java.util.HashMap; import java.util.List; import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Mono; /** * Hook for automatically forwarding agent messages to Studio. @@ -144,11 +143,12 @@ public Mono onEvent(T event) { ToolResultBlock chunk = e.getChunk(); if (chunk != null && studioClient != null) { // build a Msg with TOOL role similar to streaming hook - Msg toolMsg = Msg.builder() - .name("tool") - .role(MsgRole.TOOL) - .content(List.of(chunk)) - .build(); + Msg toolMsg = + Msg.builder() + .name("tool") + .role(MsgRole.TOOL) + .content(List.of(chunk)) + .build(); Msg tagged = withEventMetadata(toolMsg, "tool_result", false); return studioClient .pushMessage(tagged) @@ -167,11 +167,12 @@ public Mono onEvent(T event) { PostActingEvent e = (PostActingEvent) event; ToolResultBlock result = e.getToolResult(); if (result != null && studioClient != null) { - Msg toolMsg = Msg.builder() - .name("tool") - .role(MsgRole.TOOL) - .content(List.of(result)) - .build(); + Msg toolMsg = + Msg.builder() + .name("tool") + .role(MsgRole.TOOL) + .content(List.of(result)) + .build(); Msg tagged = withEventMetadata(toolMsg, "tool_result", true); return studioClient .pushMessage(tagged) @@ -198,7 +199,8 @@ private Msg withEventMetadata(Msg orig, String eventType, boolean isLast) { } meta.put("studio_event_type", eventType); meta.put("studio_is_last", isLast); - // Rebuild msg copying fields and injecting metadata; adapt if Msg.Builder has different methods + // Rebuild msg copying fields and injecting metadata; adapt if Msg.Builder has different + // methods return Msg.builder() .name(orig.getName()) .role(orig.getRole()) diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java index fa4a2361ef..4ae64efaca 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java @@ -25,24 +25,23 @@ import static org.mockito.Mockito.when; import io.agentscope.core.agent.Agent; +import io.agentscope.core.hook.ActingChunkEvent; import io.agentscope.core.hook.HookEvent; +import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.hook.PostCallEvent; +import io.agentscope.core.hook.PostReasoningEvent; import io.agentscope.core.hook.PreCallEvent; import io.agentscope.core.hook.ReasoningChunkEvent; -import io.agentscope.core.hook.PostReasoningEvent; -import io.agentscope.core.hook.ActingChunkEvent; -import io.agentscope.core.hook.PostActingEvent; -import io.agentscope.core.message.ToolResultBlock; import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolResultBlock; +import io.agentscope.core.message.ToolUseBlock; +import io.agentscope.core.tool.Toolkit; import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.Map; - -import io.agentscope.core.message.ToolUseBlock; -import io.agentscope.core.tool.Toolkit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -268,12 +267,14 @@ void testReasoningChunkEventForwards() { StepVerifier.create(result).expectNext(event).verifyComplete(); verify(mockStudioClient, times(1)) - .pushMessage(argThat(m -> { - Map md = m.getMetadata(); - return md != null - && "reasoning".equals(md.get("studio_event_type")) - && Boolean.FALSE.equals(md.get("studio_is_last")); - })); + .pushMessage( + argThat( + m -> { + Map md = m.getMetadata(); + return md != null + && "reasoning".equals(md.get("studio_event_type")) + && Boolean.FALSE.equals(md.get("studio_is_last")); + })); } @Test @@ -296,12 +297,14 @@ void testPostReasoningEventForwards() { StepVerifier.create(result).expectNext(event).verifyComplete(); verify(mockStudioClient, times(1)) - .pushMessage(argThat(m -> { - Map md = m.getMetadata(); - return md != null - && "reasoning".equals(md.get("studio_event_type")) - && Boolean.TRUE.equals(md.get("studio_is_last")); - })); + .pushMessage( + argThat( + m -> { + Map md = m.getMetadata(); + return md != null + && "reasoning".equals(md.get("studio_event_type")) + && Boolean.TRUE.equals(md.get("studio_is_last")); + })); } @Test @@ -311,7 +314,8 @@ void testActingChunkEventForwards() { when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); Toolkit toolkit = new Toolkit(); - ToolUseBlock toolUse = ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); + ToolUseBlock toolUse = + ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); ToolResultBlock chunk = ToolResultBlock.text("progress update"); ActingChunkEvent event = new ActingChunkEvent(mockAgent, toolkit, toolUse, chunk); @@ -321,15 +325,26 @@ void testActingChunkEventForwards() { StepVerifier.create(result).expectNext(event).verifyComplete(); verify(mockStudioClient, times(1)) - .pushMessage(argThat(m -> { - // Expect a TOOL role message with metadata and a ToolResultBlock in content - Map md = m.getMetadata(); - boolean hasMeta = md != null && "tool_result".equals(md.get("studio_event_type")) && Boolean.FALSE.equals(md.get("studio_is_last")); - boolean isToolRole = m.getRole() == MsgRole.TOOL; - boolean hasToolResult = - !m.getContentBlocks(io.agentscope.core.message.ToolResultBlock.class).isEmpty(); - return hasMeta && isToolRole && hasToolResult; - })); + .pushMessage( + argThat( + m -> { + // Expect a TOOL role message with metadata and a + // ToolResultBlock in content + Map md = m.getMetadata(); + boolean hasMeta = + md != null + && "tool_result" + .equals(md.get("studio_event_type")) + && Boolean.FALSE.equals( + md.get("studio_is_last")); + boolean isToolRole = m.getRole() == MsgRole.TOOL; + boolean hasToolResult = + !m.getContentBlocks( + io.agentscope.core.message + .ToolResultBlock.class) + .isEmpty(); + return hasMeta && isToolRole && hasToolResult; + })); } @Test @@ -339,7 +354,8 @@ void testPostActingEventForwards() { when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); Toolkit toolkit = new Toolkit(); - ToolUseBlock toolUse = ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); + ToolUseBlock toolUse = + ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); ToolResultBlock resultBlock = ToolResultBlock.text("final result"); PostActingEvent event = new PostActingEvent(mockAgent, toolkit, toolUse, resultBlock); @@ -349,13 +365,23 @@ void testPostActingEventForwards() { StepVerifier.create(result).expectNext(event).verifyComplete(); verify(mockStudioClient, times(1)) - .pushMessage(argThat(m -> { - Map md = m.getMetadata(); - boolean hasMeta = md != null && "tool_result".equals(md.get("studio_event_type")) && Boolean.TRUE.equals(md.get("studio_is_last")); - boolean isToolRole = m.getRole() == MsgRole.TOOL; - boolean hasToolResult = - !m.getContentBlocks(io.agentscope.core.message.ToolResultBlock.class).isEmpty(); - return hasMeta && isToolRole && hasToolResult; - })); + .pushMessage( + argThat( + m -> { + Map md = m.getMetadata(); + boolean hasMeta = + md != null + && "tool_result" + .equals(md.get("studio_event_type")) + && Boolean.TRUE.equals( + md.get("studio_is_last")); + boolean isToolRole = m.getRole() == MsgRole.TOOL; + boolean hasToolResult = + !m.getContentBlocks( + io.agentscope.core.message + .ToolResultBlock.class) + .isEmpty(); + return hasMeta && isToolRole && hasToolResult; + })); } } From 1d69a1f532f4e6c1cb909ba0333427c2fc97f4bb Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 9 Apr 2026 10:05:26 +0800 Subject: [PATCH 3/7] feature: add stream output --- .../agentscope/core/studio/StudioManager.java | 18 ++ .../core/studio/StudioMessageHook.java | 115 ------------- .../core/studio/StudioStreamingBridge.java | 155 ++++++++++++++++++ .../core/studio/StudioStreamingRunner.java | 40 +++++ .../core/studio/StudioWebSocketClient.java | 78 +++++++++ .../core/studio/StudioMessageHookTest.java | 155 ------------------ .../studio/StudioStreamingBridgeTest.java | 113 +++++++++++++ 7 files changed, 404 insertions(+), 270 deletions(-) create mode 100644 agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingBridge.java create mode 100644 agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java create mode 100644 agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java 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..67ee58c2fe 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,19 @@ public static void shutdown() { wsClient = null; } + /** + * Stream agent events to Studio. + */ + 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/StudioMessageHook.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java index 57c25e9f01..b739c2ed57 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioMessageHook.java @@ -15,19 +15,10 @@ */ package io.agentscope.core.studio; -import io.agentscope.core.hook.ActingChunkEvent; import io.agentscope.core.hook.Hook; import io.agentscope.core.hook.HookEvent; -import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.hook.PostCallEvent; -import io.agentscope.core.hook.PostReasoningEvent; -import io.agentscope.core.hook.ReasoningChunkEvent; import io.agentscope.core.message.Msg; -import io.agentscope.core.message.MsgRole; -import io.agentscope.core.message.ToolResultBlock; -import java.util.HashMap; -import java.util.List; -import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -100,112 +91,6 @@ public Mono onEvent(T event) { return Mono.just(event); }); } - - // Reasoning incremental chunk - if (event instanceof ReasoningChunkEvent) { - ReasoningChunkEvent e = (ReasoningChunkEvent) event; - Msg chunk = e.getIncrementalChunk(); - if (chunk != null && studioClient != null) { - Msg tagged = withEventMetadata(chunk, "reasoning", false); - return studioClient - .pushMessage(tagged) - .thenReturn(event) - .onErrorResume( - ex -> { - logger.error("Failed to push reasoning chunk to Studio", ex); - return Mono.just(event); - }); - } - return Mono.just(event); - } - - // Reasoning final result (after stream completes) - if (event instanceof PostReasoningEvent) { - PostReasoningEvent e = (PostReasoningEvent) event; - Msg finalMsg = e.getReasoningMessage(); - if (finalMsg != null && studioClient != null) { - Msg tagged = withEventMetadata(finalMsg, "reasoning", true); - return studioClient - .pushMessage(tagged) - .thenReturn(event) - .onErrorResume( - ex -> { - logger.error("Failed to push final reasoning to Studio", ex); - return Mono.just(event); - }); - } - return Mono.just(event); - } - - // Acting (tool) incremental chunk - if (event instanceof ActingChunkEvent) { - ActingChunkEvent e = (ActingChunkEvent) event; - ToolResultBlock chunk = e.getChunk(); - if (chunk != null && studioClient != null) { - // build a Msg with TOOL role similar to streaming hook - Msg toolMsg = - Msg.builder() - .name("tool") - .role(MsgRole.TOOL) - .content(List.of(chunk)) - .build(); - Msg tagged = withEventMetadata(toolMsg, "tool_result", false); - return studioClient - .pushMessage(tagged) - .thenReturn(event) - .onErrorResume( - ex -> { - logger.error("Failed to push acting chunk to Studio", ex); - return Mono.just(event); - }); - } - return Mono.just(event); - } - - // Acting (tool) final result - if (event instanceof PostActingEvent) { - PostActingEvent e = (PostActingEvent) event; - ToolResultBlock result = e.getToolResult(); - if (result != null && studioClient != null) { - Msg toolMsg = - Msg.builder() - .name("tool") - .role(MsgRole.TOOL) - .content(List.of(result)) - .build(); - Msg tagged = withEventMetadata(toolMsg, "tool_result", true); - return studioClient - .pushMessage(tagged) - .thenReturn(event) - .onErrorResume( - ex -> { - logger.error("Failed to push tool result to Studio", ex); - return Mono.just(event); - }); - } - return Mono.just(event); - } - return Mono.just(event); } - - /** - * Wrap/augment a Msg with metadata that indicates event type and whether it's the last chunk. - */ - private Msg withEventMetadata(Msg orig, String eventType, boolean isLast) { - Map meta = new HashMap<>(); - if (orig.getMetadata() != null) { - meta.putAll(orig.getMetadata()); - } - meta.put("studio_event_type", eventType); - meta.put("studio_is_last", isLast); - // Rebuild msg copying fields and injecting metadata; adapt if Msg.Builder has different - // methods - return Msg.builder() - .name(orig.getName()) - .role(orig.getRole()) - .content(orig.getContent()) - .metadata(meta) - .build(); - } } 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..069760c095 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingBridge.java @@ -0,0 +1,155 @@ +/* + * 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 org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.util.Objects; + +public class StudioStreamingBridge { + + private static final Logger logger = LoggerFactory.getLogger(StudioStreamingBridge.class); + + private final StudioClient studioClient; + private final StudioWebSocketClient webSocketClient; + + 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 Flux to Studio: + * + *
    + *
  1. Non-terminal events (isLast == false) are streamed to the frontend via + * webSocketClient.sendStreamEvent(event).
    2. + *
  2. The terminal AGENT_RESULT event with isLast == true is used to: + *
      + *
    • cache the final Msg for persistence, and
      • + *
    • trigger webSocketClient.sendStreamCompleted().
      • + *
    + * It is not sent again as a streaming chunk.
    3. + *
  3. If a final Msg exists, it is persisted once via studioClient.pushMessage(finalMsg) + * after the stream completes.
    4. + *
+ * + *

Note: isLast is treated purely as a control flag 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 (!isLast + && (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/StudioStreamingRunner.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java new file mode 100644 index 0000000000..19b9003c7c --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java @@ -0,0 +1,40 @@ +/* + * 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.StreamOptions; +import io.agentscope.core.agent.StreamableAgent; +import io.agentscope.core.message.Msg; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +public class StudioStreamingRunner { + + private final StudioStreamingBridge streamingBridge; + + public StudioStreamingRunner(StudioStreamingBridge streamingBridge) { + this.streamingBridge = streamingBridge; + } + + public Mono runWithStreaming(StreamableAgent agent, Msg input, StreamOptions options) { + StreamOptions effectiveOptions = + options != null ? options : StreamOptions.defaults(); + + Flux eventFlux = agent.stream(input, effectiveOptions); + return streamingBridge.forwardToStudio(eventFlux); + } +} 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..b7f3d2586d 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,11 @@ */ 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.util.JsonUtils; import io.socket.client.IO; import io.socket.client.Socket; @@ -395,4 +399,78 @@ public Map getStructuredInput() { return structuredInput; } } + + /** + * Sends a stream event to Studio in real-time. + */ + 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. + */ + 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; + } + Object first = msg.getContent().get(0); + if (first instanceof TextBlock tb && tb.getText() != null) { + return tb.getText(); + } + + return null; + } } diff --git a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java index 4ae64efaca..48b07bc8b8 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioMessageHookTest.java @@ -17,7 +17,6 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -25,23 +24,15 @@ import static org.mockito.Mockito.when; import io.agentscope.core.agent.Agent; -import io.agentscope.core.hook.ActingChunkEvent; import io.agentscope.core.hook.HookEvent; -import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.hook.PostCallEvent; -import io.agentscope.core.hook.PostReasoningEvent; import io.agentscope.core.hook.PreCallEvent; -import io.agentscope.core.hook.ReasoningChunkEvent; import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; -import io.agentscope.core.message.ToolResultBlock; -import io.agentscope.core.message.ToolUseBlock; -import io.agentscope.core.tool.Toolkit; import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.Map; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -238,150 +229,4 @@ void testMultipleEvents() { verify(mockStudioClient, times(1)).pushMessage(msg1); verify(mockStudioClient, times(1)).pushMessage(msg2); } - - @Test - @DisplayName("Should forward reasoning chunk to Studio") - void testReasoningChunkEventForwards() { - // Mock push success - when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); - - // Build incremental and accumulated msg - Msg incremental = - Msg.builder() - .name("Assistant") - .role(MsgRole.ASSISTANT) - .content(TextBlock.builder().text("partial chunk").build()) - .build(); - Msg accumulated = - Msg.builder() - .name("Assistant") - .role(MsgRole.ASSISTANT) - .content(TextBlock.builder().text("accumulated so far").build()) - .build(); - - ReasoningChunkEvent event = - new ReasoningChunkEvent(mockAgent, "test-model", null, incremental, accumulated); - - Mono result = hook.onEvent(event); - - StepVerifier.create(result).expectNext(event).verifyComplete(); - - verify(mockStudioClient, times(1)) - .pushMessage( - argThat( - m -> { - Map md = m.getMetadata(); - return md != null - && "reasoning".equals(md.get("studio_event_type")) - && Boolean.FALSE.equals(md.get("studio_is_last")); - })); - } - - @Test - @DisplayName("Should forward final reasoning result to Studio") - void testPostReasoningEventForwards() { - // Mock push success - when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); - - Msg finalMsg = - Msg.builder() - .name("Assistant") - .role(MsgRole.ASSISTANT) - .content(TextBlock.builder().text("final reasoning").build()) - .build(); - - PostReasoningEvent event = new PostReasoningEvent(mockAgent, "test-model", null, finalMsg); - - Mono result = hook.onEvent(event); - - StepVerifier.create(result).expectNext(event).verifyComplete(); - - verify(mockStudioClient, times(1)) - .pushMessage( - argThat( - m -> { - Map md = m.getMetadata(); - return md != null - && "reasoning".equals(md.get("studio_event_type")) - && Boolean.TRUE.equals(md.get("studio_is_last")); - })); - } - - @Test - @DisplayName("Should forward acting (tool) chunk to Studio") - void testActingChunkEventForwards() { - // Mock push success - when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); - - Toolkit toolkit = new Toolkit(); - ToolUseBlock toolUse = - ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); - ToolResultBlock chunk = ToolResultBlock.text("progress update"); - - ActingChunkEvent event = new ActingChunkEvent(mockAgent, toolkit, toolUse, chunk); - - Mono result = hook.onEvent(event); - - StepVerifier.create(result).expectNext(event).verifyComplete(); - - verify(mockStudioClient, times(1)) - .pushMessage( - argThat( - m -> { - // Expect a TOOL role message with metadata and a - // ToolResultBlock in content - Map md = m.getMetadata(); - boolean hasMeta = - md != null - && "tool_result" - .equals(md.get("studio_event_type")) - && Boolean.FALSE.equals( - md.get("studio_is_last")); - boolean isToolRole = m.getRole() == MsgRole.TOOL; - boolean hasToolResult = - !m.getContentBlocks( - io.agentscope.core.message - .ToolResultBlock.class) - .isEmpty(); - return hasMeta && isToolRole && hasToolResult; - })); - } - - @Test - @DisplayName("Should forward final acting (tool) result to Studio") - void testPostActingEventForwards() { - // Mock push success - when(mockStudioClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); - - Toolkit toolkit = new Toolkit(); - ToolUseBlock toolUse = - ToolUseBlock.builder().id("call-1").name("test_tool").input(Map.of()).build(); - ToolResultBlock resultBlock = ToolResultBlock.text("final result"); - - PostActingEvent event = new PostActingEvent(mockAgent, toolkit, toolUse, resultBlock); - - Mono result = hook.onEvent(event); - - StepVerifier.create(result).expectNext(event).verifyComplete(); - - verify(mockStudioClient, times(1)) - .pushMessage( - argThat( - m -> { - Map md = m.getMetadata(); - boolean hasMeta = - md != null - && "tool_result" - .equals(md.get("studio_event_type")) - && Boolean.TRUE.equals( - md.get("studio_is_last")); - boolean isToolRole = m.getRole() == MsgRole.TOOL; - boolean hasToolResult = - !m.getContentBlocks( - io.agentscope.core.message - .ToolResultBlock.class) - .isEmpty(); - return hasMeta && isToolRole && hasToolResult; - })); - } } 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..42341eae4d --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java @@ -0,0 +1,113 @@ +/* + * 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.Mockito.*; + +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 java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +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(); + } + + @Test + void testStreamingAndCompletion() { + + StudioClient mockClient = mock(StudioClient.class); + StudioWebSocketClient mockWs = 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, 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 = mock(StudioClient.class); + StudioWebSocketClient mockWs = 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)); + } +} + From 7b9c0c6312392bed3312da0ec7de0250170ea5fe Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 9 Apr 2026 10:58:11 +0800 Subject: [PATCH 4/7] fix ci --- .../agentscope/core/studio/StudioManager.java | 17 +++- .../core/studio/StudioStreamingBridge.java | 71 +++++++++++----- .../core/studio/StudioStreamingRunner.java | 40 --------- .../core/studio/StudioWebSocketClient.java | 59 ++++++++++++- .../studio/StudioStreamingBridgeTest.java | 82 ++++++++++++++++--- 5 files changed, 191 insertions(+), 78 deletions(-) delete mode 100644 agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java 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 67ee58c2fe..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 @@ -127,9 +127,22 @@ public static void shutdown() { } /** - * Stream agent events to Studio. + * 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) { + 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")); } 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 index 069760c095..0dd953c304 100644 --- 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 @@ -15,17 +15,35 @@ */ 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; -import java.util.Objects; - +/** + * 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); @@ -33,6 +51,13 @@ public class StudioStreamingBridge { 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 = @@ -40,23 +65,28 @@ public StudioStreamingBridge(StudioClient studioClient, StudioWebSocketClient we } /** - * Forwards Flux to Studio: + * Forwards a stream of agent events to Studio for real-time visualization and + * final message persistence. * *

    - *
  1. Non-terminal events (isLast == false) are streamed to the frontend via - * webSocketClient.sendStreamEvent(event).
    2. - *
  2. The terminal AGENT_RESULT event with isLast == true is used to: + *
  3. 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.
    4. + *
  4. Non-terminal AGENT_RESULT events ({@code isLast == false}) are also forwarded as + * streaming chunks.
    5. + *
  5. The terminal AGENT_RESULT event with {@code isLast == true} is used to: *
      - *
    • cache the final Msg for persistence, and
      • - *
    • trigger webSocketClient.sendStreamCompleted().
      • + *
    • cache the final {@link Msg} for persistence, and
      • + *
    • trigger {@link StudioWebSocketClient#sendStreamCompleted()}.
      • *
    * It is not sent again as a streaming chunk.
    6. - *
  6. If a final Msg exists, it is persisted once via studioClient.pushMessage(finalMsg) - * after the stream completes.
    7. + *
  7. If a final {@link Msg} exists, it is persisted once via + * {@link StudioClient#pushMessage(Msg)} after the stream completes.
    8. *
* - *

Note: isLast is treated purely as a control flag and does not generate - * additional visible content.

+ *

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( @@ -68,7 +98,8 @@ public Mono forwardToStudio(Flux eventFlux) { .doOnError( ex -> logger.error( - "Error occurred during streaming to Studio", ex)) + "Error occurred during streaming to Studio", + ex)) .doFinally( signalType -> { if (!holder.completedSignalSent) { @@ -76,7 +107,8 @@ public Mono forwardToStudio(Flux eventFlux) { webSocketClient.sendStreamCompleted(); } catch (Exception e) { logger.error( - "Failed to send streamCompleted in doFinally", + "Failed to send streamCompleted in" + + " doFinally", e); } holder.completedSignalSent = true; @@ -88,8 +120,8 @@ public Mono forwardToStudio(Flux eventFlux) { Msg finalMsg = holder.getEffectiveFinalMsg(); if (finalMsg == null) { logger.debug( - "No final Msg determined from stream; skip" - + " pushMessage"); + "No final Msg determined from stream;" + + " skip pushMessage"); return Mono.empty(); } return studioClient.pushMessage(finalMsg); @@ -129,10 +161,9 @@ private void handleEvent(Event event, FinalMsgHolder holder) { return; } - if (!isLast - && (type == EventType.REASONING + if (type == EventType.REASONING || type == EventType.TOOL_RESULT - || type == EventType.SUMMARY)) { + || type == EventType.SUMMARY) { webSocketClient.sendStreamEvent(event); } } diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java deleted file mode 100644 index 19b9003c7c..0000000000 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioStreamingRunner.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * 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.StreamOptions; -import io.agentscope.core.agent.StreamableAgent; -import io.agentscope.core.message.Msg; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -public class StudioStreamingRunner { - - private final StudioStreamingBridge streamingBridge; - - public StudioStreamingRunner(StudioStreamingBridge streamingBridge) { - this.streamingBridge = streamingBridge; - } - - public Mono runWithStreaming(StreamableAgent agent, Msg input, StreamOptions options) { - StreamOptions effectiveOptions = - options != null ? options : StreamOptions.defaults(); - - Flux eventFlux = agent.stream(input, effectiveOptions); - return streamingBridge.forwardToStudio(eventFlux); - } -} 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 b7f3d2586d..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 @@ -20,6 +20,8 @@ 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; @@ -402,6 +404,14 @@ public Map getStructuredInput() { /** * 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()) { @@ -444,6 +454,11 @@ public void sendStreamEvent(Event event) { /** * 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()) { @@ -466,11 +481,47 @@ private String extractTextFromMsg(Msg msg) { if (msg == null || msg.getContent() == null || msg.getContent().isEmpty()) { return null; } - Object first = msg.getContent().get(0); - if (first instanceof TextBlock tb && tb.getText() != null) { - return tb.getText(); + + 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 null; + 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/StudioStreamingBridgeTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioStreamingBridgeTest.java index 42341eae4d..aef9c8a3da 100644 --- 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 @@ -16,17 +16,23 @@ package io.agentscope.core.studio; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.*; +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; @@ -42,11 +48,29 @@ private Msg msg(String text) { .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 = mock(StudioClient.class); - StudioWebSocketClient mockWs = mock(StudioWebSocketClient.class); + StudioClient mockClient = Mockito.mock(StudioClient.class); + StudioWebSocketClient mockWs = Mockito.mock(StudioWebSocketClient.class); when(mockClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); @@ -77,11 +101,14 @@ void testStreamingAndCompletion() { verify(mockWs, times(5)).sendStreamEvent(any(Event.class)); - verify(mockWs, atLeastOnce()).sendStreamEvent(argThat(e -> - !e.isLast() - && (e.getType() == EventType.REASONING - || e.getType() == EventType.TOOL_RESULT - || e.getType() == EventType.AGENT_RESULT))); + 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(); @@ -90,8 +117,8 @@ void testStreamingAndCompletion() { @Test void testFallbackWhenNoTerminalAgentResult() { - StudioClient mockClient = mock(StudioClient.class); - StudioWebSocketClient mockWs = mock(StudioWebSocketClient.class); + StudioClient mockClient = Mockito.mock(StudioClient.class); + StudioWebSocketClient mockWs = Mockito.mock(StudioWebSocketClient.class); when(mockClient.pushMessage(any(Msg.class))).thenReturn(Mono.empty()); @@ -109,5 +136,36 @@ void testFallbackWhenNoTerminalAgentResult() { 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)); + } +} From b2a61bdd646e384ea382be29bd14f1d962a23690 Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 9 Apr 2026 11:46:32 +0800 Subject: [PATCH 5/7] fomate the code style --- .../core/studio/StudioStreamingBridgeTest.java | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) 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 index aef9c8a3da..fe3dbe80d9 100644 --- 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 @@ -57,13 +57,8 @@ private Msg thinkingMsg(String thinking) { } private Msg toolResultMsg(String text) { - ToolResultBlock toolResult = - ToolResultBlock.text(text); - return Msg.builder() - .role(MsgRole.ASSISTANT) - .name("Agent") - .content(toolResult) - .build(); + ToolResultBlock toolResult = ToolResultBlock.text(text); + return Msg.builder().role(MsgRole.ASSISTANT).name("Agent").content(toolResult).build(); } @Test @@ -148,8 +143,7 @@ void testReasoningAndToolResultEventsWithThinkingAndToolResultBlocksAreStreamed( Event reasoningFinal = new Event(EventType.REASONING, thinkingMsg("chain-of-thought"), true); - Event toolFinal = - new Event(EventType.TOOL_RESULT, toolResultMsg("tool-output"), 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"); From 1faf1978cc131ef802e10107e686b8fadfffce70 Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 9 Apr 2026 12:49:05 +0800 Subject: [PATCH 6/7] add some tests --- .../core/studio/StudioManagerTest.java | 105 +++++- .../studio/StudioWebSocketClientTest.java | 330 +++++++++++++++++- 2 files changed, 433 insertions(+), 2 deletions(-) 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/StudioWebSocketClientTest.java b/agentscope-extensions/agentscope-extensions-studio/src/test/java/io/agentscope/core/studio/StudioWebSocketClientTest.java index 9830cb8cc2..98f55c0b16 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,12 +18,30 @@ 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.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 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 reactor.test.StepVerifier; @DisplayName("StudioWebSocketClient Tests") @@ -304,4 +322,314 @@ void testConstructorWithSocket() { assertEquals(true, clientWithSocket.isConnected()); clientWithSocket.close(); } + + // ==================== sendStreamEvent Tests ==================== + + @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(); + } } From 05dac8f7b73bb4c91c430af4d22190c3051a3159 Mon Sep 17 00:00:00 2001 From: Loki Date: Thu, 9 Apr 2026 13:07:34 +0800 Subject: [PATCH 7/7] formate code style --- .../core/studio/StudioWebSocketClientTest.java | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) 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 f00eacd071..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 @@ -28,6 +28,7 @@ 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; @@ -35,15 +36,10 @@ import io.agentscope.core.message.ToolResultBlock; import io.socket.client.Socket; import java.time.Duration; -import org.json.JSONException; - -import io.agentscope.core.message.ContentBlock; -import io.agentscope.core.message.TextBlock; -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; @@ -324,8 +320,6 @@ void testConstructorWithSocket() { clientWithSocket.close(); } - // ==================== sendStreamEvent Tests ==================== - @Test @DisplayName("sendStreamEvent should log warning when socket is null") void sendStreamEvent_WithNullSocket_ShouldLogWarning() {