From 00d825db9e92b74acda8c896fb933130e22093ed Mon Sep 17 00:00:00 2001 From: SamizuHM Date: Fri, 31 Jul 2026 04:07:31 +0800 Subject: [PATCH] fix(stream): preserve terminal events and completion --- .../client/callback/BaseStreamCallback.java | 9 ++ .../dify/client/impl/AbstractDifyClient.java | 95 ++++++++++++- .../dify/client/impl/DefaultDifyClient.java | 63 +++++++-- .../impl/DefaultDifyDatasetsClient.java | 12 +- .../ChatflowStreamTerminalEventTest.java | 131 ++++++++++++++++++ .../client/PipelineStreamCompletionTest.java | 83 +++++++++++ .../WorkflowEventsStreamCompletionTest.java | 72 ++++++++++ 7 files changed, 441 insertions(+), 24 deletions(-) create mode 100644 src/test/java/io/github/imfangs/dify/client/ChatflowStreamTerminalEventTest.java create mode 100644 src/test/java/io/github/imfangs/dify/client/PipelineStreamCompletionTest.java create mode 100644 src/test/java/io/github/imfangs/dify/client/WorkflowEventsStreamCompletionTest.java diff --git a/src/main/java/io/github/imfangs/dify/client/callback/BaseStreamCallback.java b/src/main/java/io/github/imfangs/dify/client/callback/BaseStreamCallback.java index eddf98e..bd94171 100644 --- a/src/main/java/io/github/imfangs/dify/client/callback/BaseStreamCallback.java +++ b/src/main/java/io/github/imfangs/dify/client/callback/BaseStreamCallback.java @@ -33,4 +33,13 @@ default void onPing(PingEvent event) { default void onException(Throwable throwable) { } + /** + * SDK 已完成对 SSE 响应的读取。 + * + *

当收到客户端定义的终止事件，或服务端正常关闭 SSE 连接时触发。该回调仅表示协议流读取结束， + * 不表示 Dify 工作流或调用方业务处理成功。Dify {@code error} 事件和网络异常不会触发该回调。

+ */ + default void onStreamComplete() { + } + } diff --git a/src/main/java/io/github/imfangs/dify/client/impl/AbstractDifyClient.java b/src/main/java/io/github/imfangs/dify/client/impl/AbstractDifyClient.java index f77c3c3..caa5bba 100644 --- a/src/main/java/io/github/imfangs/dify/client/impl/AbstractDifyClient.java +++ b/src/main/java/io/github/imfangs/dify/client/impl/AbstractDifyClient.java @@ -491,6 +491,23 @@ protected interface LineProcessor { boolean process(String line); } + /** + * 带结束原因的行处理器。 + */ + @FunctionalInterface + protected interface StreamLineProcessor { + StreamLineResult process(String line); + } + + /** + * SSE 行处理结果。 + */ + protected enum StreamLineResult { + CONTINUE, + COMPLETE, + ERROR + } + /** * 事件处理器 */ @@ -514,6 +531,25 @@ protected void executeStreamRequest(String path, Object body, LineProcessor line executeStreamCall(httpRequest, lineProcessor, errorHandler); } + /** + * 执行 POST 流式请求，并在正常结束时通知调用方。 + */ + protected void executeStreamRequest(String path, + Object body, + StreamLineProcessor lineProcessor, + Consumer errorHandler, + Runnable completionHandler) { + RequestBody requestBody = createJsonRequestBody(body); + Request httpRequest = new Request.Builder() + .url(baseUrl + path) + .post(requestBody) + .header("Authorization", "Bearer " + apiKey) + .header("Content-Type", "application/json") + .header("Accept", "text/event-stream") + .build(); + executeStreamCall(httpRequest, lineProcessor, errorHandler, completionHandler); + } + /** * 执行 GET 流式请求 */ @@ -527,7 +563,33 @@ protected void executeGetStreamRequest(String path, LineProcessor lineProcessor, executeStreamCall(httpRequest, lineProcessor, errorHandler); } + /** + * 执行 GET 流式请求，并在正常结束时通知调用方。 + */ + protected void executeGetStreamRequest(String path, + StreamLineProcessor lineProcessor, + Consumer errorHandler, + Runnable completionHandler) { + Request httpRequest = new Request.Builder() + .url(baseUrl + path) + .get() + .header("Authorization", "Bearer " + apiKey) + .header("Accept", "text/event-stream") + .build(); + executeStreamCall(httpRequest, lineProcessor, errorHandler, completionHandler); + } + protected void executeStreamCall(Request httpRequest, LineProcessor lineProcessor, Consumer errorHandler) { + executeStreamCall(httpRequest, + line -> lineProcessor.process(line) ? StreamLineResult.CONTINUE : StreamLineResult.COMPLETE, + errorHandler, + null); + } + + protected void executeStreamCall(Request httpRequest, + StreamLineProcessor lineProcessor, + Consumer errorHandler, + Runnable completionHandler) { Call call = httpClient.newCall(httpRequest); call.enqueue(new Callback() { @Override @@ -561,13 +623,23 @@ public void onResponse(Call call, Response response) { try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream(), StandardCharsets.UTF_8))) { String line; + boolean receivedErrorEvent = false; while ((line = reader.readLine()) != null) { if (line.isEmpty()) { continue; } - if (!lineProcessor.process(line)) { + + StreamLineResult result = lineProcessor.process(line); + if (result == StreamLineResult.ERROR) { + receivedErrorEvent = true; break; } + if (result == StreamLineResult.COMPLETE) { + break; + } + } + if (!receivedErrorEvent && completionHandler != null) { + completionHandler.run(); } } } catch (Exception e) { @@ -582,8 +654,18 @@ public void onResponse(Call call, Response response) { * 处理一行 SSE 数据 */ protected boolean processStreamLine(String line, BaseStreamCallback callback, Set terminalEvents, EventProcessor eventProcessor) { + return processStreamLineWithResult(line, callback, terminalEvents, eventProcessor) == StreamLineResult.CONTINUE; + } + + /** + * 处理一行 SSE 数据，并保留结束原因。 + */ + protected StreamLineResult processStreamLineWithResult(String line, + BaseStreamCallback callback, + Set terminalEvents, + EventProcessor eventProcessor) { if (line == null || line.trim().isEmpty()) { - return true; + return StreamLineResult.CONTINUE; } if (line.startsWith(STREAM_DATA_PREFIX)) { String data = line.substring(STREAM_DATA_PREFIX.length()).trim(); @@ -591,13 +673,16 @@ protected boolean processStreamLine(String line, BaseStreamCallback callback, Se BaseEvent baseEvent = JsonUtils.fromJson(data, BaseEvent.class); if (baseEvent == null) { log.warn("解析事件数据为null: {}", data); - return true; + return StreamLineResult.CONTINUE; } eventProcessor.process(data, baseEvent.getEvent()); String eventTypeStr = baseEvent.getEvent(); EventType eventType = eventTypeStr != null ? EventType.fromValue(eventTypeStr) : null; + if (eventType == EventType.ERROR) { + return StreamLineResult.ERROR; + } if (eventType != null && terminalEvents.contains(eventType)) { - return false; + return StreamLineResult.COMPLETE; } } catch (Exception e) { log.error("解析事件数据失败: {}", data, e); @@ -608,6 +693,6 @@ protected boolean processStreamLine(String line, BaseStreamCallback callback, Se pingEvent.setEvent(EventType.PING.getValue()); callback.onPing(pingEvent); } - return true; + return StreamLineResult.CONTINUE; } } diff --git a/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyClient.java b/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyClient.java index 0d97f9d..d4073a0 100644 --- a/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyClient.java +++ b/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyClient.java @@ -18,6 +18,7 @@ import java.io.*; import java.nio.charset.StandardCharsets; +import java.util.Collections; import java.util.EnumSet; import java.util.HashMap; import java.util.Map; @@ -33,8 +34,9 @@ public class DefaultDifyClient extends DifyBaseClientImpl implements DifyClient { // 流式响应相关常量 - private static final Set CHAT_TERMINAL_EVENTS = EnumSet.of(EventType.MESSAGE_END, EventType.ERROR); - private static final Set WORKFLOW_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED, EventType.ERROR); + private static final Set CHAT_TERMINAL_EVENTS = EnumSet.of(EventType.MESSAGE_END); + private static final Set NO_TERMINAL_EVENTS = Collections.emptySet(); + private static final Set WORKFLOW_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED); // API 路径常量 // 对话型应用相关路径 @@ -111,9 +113,9 @@ public void sendChatMessageStream(ChatMessage message, ChatStreamCallback callba message.setResponseMode(ResponseMode.STREAMING); // 执行流式请求 - executeStreamRequest(CHAT_MESSAGES_PATH, message, (line) -> processStreamLine(line, callback, CHAT_TERMINAL_EVENTS, (data, eventType) -> { + executeStreamRequest(CHAT_MESSAGES_PATH, message, (line) -> processStreamLineWithResult(line, callback, CHAT_TERMINAL_EVENTS, (data, eventType) -> { StreamEventDispatcher.dispatchChatEvent(callback, data, eventType); - }), callback::onException); + }), callback::onException, callback::onStreamComplete); } @Override @@ -122,10 +124,19 @@ public void sendChatMessageStream(ChatMessage message, ChatflowStreamCallback ca // 确保请求模式为流式 message.setResponseMode(ResponseMode.STREAMING); - // 执行流式请求 - executeStreamRequest(CHAT_MESSAGES_PATH, message, (line) -> processStreamLine(line, callback, WORKFLOW_TERMINAL_EVENTS, (data, eventType) -> { - StreamEventDispatcher.dispatchChatFlowEvent(callback, data, eventType); - }), callback::onException); + ChatflowTerminalState terminalState = new ChatflowTerminalState(); + executeStreamRequest(CHAT_MESSAGES_PATH, message, (line) -> { + StreamLineResult result = processStreamLineWithResult(line, callback, NO_TERMINAL_EVENTS, (data, eventType) -> { + StreamEventDispatcher.dispatchChatFlowEvent(callback, data, eventType); + terminalState.record(eventType); + }); + if (result != StreamLineResult.CONTINUE) { + return result; + } + return terminalState.hasReceivedAllTerminalEvents() + ? StreamLineResult.COMPLETE + : StreamLineResult.CONTINUE; + }, callback::onException, callback::onStreamComplete); } @Override @@ -290,10 +301,10 @@ public void sendCompletionMessageStream(CompletionRequest request, CompletionStr request.setResponseMode(ResponseMode.STREAMING); // 执行流式请求 - executeStreamRequest(COMPLETION_MESSAGES_PATH, request, (line) -> processStreamLine(line, callback, CHAT_TERMINAL_EVENTS, (data, eventType) -> { + executeStreamRequest(COMPLETION_MESSAGES_PATH, request, (line) -> processStreamLineWithResult(line, callback, CHAT_TERMINAL_EVENTS, (data, eventType) -> { // 分发事件 StreamEventDispatcher.dispatchCompletionEvent(callback, data); - }), callback::onException); + }), callback::onException, callback::onStreamComplete); } @Override @@ -326,10 +337,10 @@ public void runWorkflowStream(WorkflowRunRequest request, WorkflowStreamCallback request.setResponseMode(ResponseMode.STREAMING); // 执行流式请求 - executeStreamRequest(WORKFLOWS_RUN_PATH, request, (line) -> processStreamLine(line, callback, WORKFLOW_TERMINAL_EVENTS, (data, eventType) -> { + executeStreamRequest(WORKFLOWS_RUN_PATH, request, (line) -> processStreamLineWithResult(line, callback, WORKFLOW_TERMINAL_EVENTS, (data, eventType) -> { // 分发事件 StreamEventDispatcher.dispatchWorkflowEvent(callback, data); - }), callback::onException); + }), callback::onException, callback::onStreamComplete); } @Override @@ -408,9 +419,9 @@ public void streamWorkflowEvents(String workflowRunId, params.put("continue_on_pause", continueOnPause); } String url = buildUrlWithParams(WORKFLOW_EVENTS_PATH + "/" + workflowRunId.trim() + "/events", params); - executeGetStreamRequest(url, (line) -> processStreamLine(line, callback, WORKFLOW_TERMINAL_EVENTS, (data, eventType) -> { + executeGetStreamRequest(url, (line) -> processStreamLineWithResult(line, callback, WORKFLOW_TERMINAL_EVENTS, (data, eventType) -> { StreamEventDispatcher.dispatchWorkflowEvent(callback, data); - }), callback::onException); + }), callback::onException, callback::onStreamComplete); } /** @@ -554,4 +565,28 @@ private String getFileExtension(String fileName) { int lastDotIndex = fileName.lastIndexOf('.'); return lastDotIndex > 0 ? fileName.substring(lastDotIndex + 1) : ""; } + + /** + * 跟踪 Chatflow 的两个终止事件。 + * + *

不同 Dify 版本可能以不同顺序发送 {@code message_end} 和 + * {@code workflow_finished}。只有两者都收到时，客户端才主动停止读取；若连接先关闭， + * 则由通用的流结束回调处理。

+ */ + private static final class ChatflowTerminalState { + private boolean messageEndReceived; + private boolean workflowFinishedReceived; + + private void record(String eventType) { + if (EventType.MESSAGE_END.getValue().equals(eventType)) { + messageEndReceived = true; + } else if (EventType.WORKFLOW_FINISHED.getValue().equals(eventType)) { + workflowFinishedReceived = true; + } + } + + private boolean hasReceivedAllTerminalEvents() { + return messageEndReceived && workflowFinishedReceived; + } + } } diff --git a/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyDatasetsClient.java b/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyDatasetsClient.java index e137573..aaa4ab6 100644 --- a/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyDatasetsClient.java +++ b/src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyDatasetsClient.java @@ -635,7 +635,7 @@ public DocumentDownloadUrlResponse getDocumentDownloadUrl(String datasetId, Stri private static final String PIPELINE_DATASOURCE_NODES_PATH = "/pipeline/datasource/nodes"; private static final String PIPELINE_RUN_PATH = "/pipeline/run"; private static final String PIPELINE_FILE_UPLOAD_PATH = "/datasets/pipeline/file-upload"; - private static final Set PIPELINE_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED, EventType.ERROR); + private static final Set PIPELINE_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED); @Override public List listPipelineDatasourcePlugins(String datasetId, Boolean isPublished) throws IOException, DifyApiException { @@ -676,9 +676,10 @@ public void runPipelineDatasourceNodeStream(String datasetId, log.debug("运行 Pipeline 数据源节点: datasetId={}, nodeId={}, request={}", datasetId, nodeId, request); String path = DATASETS_PATH + "/" + datasetId + PIPELINE_DATASOURCE_NODES_PATH + "/" + nodeId + "/run"; executeStreamRequest(path, request, - (line) -> processStreamLine(line, callback, PIPELINE_TERMINAL_EVENTS, + (line) -> processStreamLineWithResult(line, callback, PIPELINE_TERMINAL_EVENTS, (data, eventType) -> StreamEventDispatcher.dispatchWorkflowEvent(callback, data)), - callback::onException); + callback::onException, + callback::onStreamComplete); } @Override @@ -719,9 +720,10 @@ public void runPipelineStream(String datasetId, log.debug("运行 Pipeline（流式模式）: datasetId={}", datasetId); String path = DATASETS_PATH + "/" + datasetId + PIPELINE_RUN_PATH; executeStreamRequest(path, request, - (line) -> processStreamLine(line, callback, PIPELINE_TERMINAL_EVENTS, + (line) -> processStreamLineWithResult(line, callback, PIPELINE_TERMINAL_EVENTS, (data, eventType) -> StreamEventDispatcher.dispatchWorkflowEvent(callback, data)), - callback::onException); + callback::onException, + callback::onStreamComplete); } @Override diff --git a/src/test/java/io/github/imfangs/dify/client/ChatflowStreamTerminalEventTest.java b/src/test/java/io/github/imfangs/dify/client/ChatflowStreamTerminalEventTest.java new file mode 100644 index 0000000..3dc060e --- /dev/null +++ b/src/test/java/io/github/imfangs/dify/client/ChatflowStreamTerminalEventTest.java @@ -0,0 +1,131 @@ +package io.github.imfangs.dify.client; + +import io.github.imfangs.dify.client.callback.ChatflowStreamCallback; +import io.github.imfangs.dify.client.event.ErrorEvent; +import io.github.imfangs.dify.client.event.MessageEndEvent; +import io.github.imfangs.dify.client.event.WorkflowFinishedEvent; +import io.github.imfangs.dify.client.impl.DefaultDifyClient; +import io.github.imfangs.dify.client.model.chat.ChatMessage; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * 验证 Chatflow 的终止事件顺序和连接关闭都能正确结束流处理。 + */ +class ChatflowStreamTerminalEventTest { + + @Test + void shouldDeliverBothTerminalEventsInDocumentedOrder() throws Exception { + assertTerminalEvents( + "data: {\"event\":\"message_end\",\"message_id\":\"message-1\",\"conversation_id\":\"conversation-1\",\"metadata\":{}}\n\n" + + "data: {\"event\":\"workflow_finished\",\"data\":{\"id\":\"workflow-1\",\"status\":\"succeeded\"}}\n\n", + Arrays.asList("message_end", "workflow_finished", "stream_complete")); + } + + @Test + void shouldDeliverBothTerminalEventsInReverseOrder() throws Exception { + assertTerminalEvents( + "data: {\"event\":\"workflow_finished\",\"data\":{\"id\":\"workflow-1\",\"status\":\"succeeded\"}}\n\n" + + "data: {\"event\":\"message_end\",\"message_id\":\"message-1\",\"conversation_id\":\"conversation-1\",\"metadata\":{}}\n\n", + Arrays.asList("workflow_finished", "message_end", "stream_complete")); + } + + @Test + void shouldNotifyCompletionWhenConnectionClosesBeforeAllTerminalEventsArrive() throws Exception { + CountDownLatch streamCompleted = new CountDownLatch(1); + DefaultDifyClient client = createClient( + "data: {\"event\":\"workflow_finished\",\"data\":{\"id\":\"workflow-1\",\"status\":\"succeeded\"}}\n\n"); + + client.sendChatMessageStream(ChatMessage.builder().query("test").user("tester").build(), + new ChatflowStreamCallback() { + @Override + public void onStreamComplete() { + streamCompleted.countDown(); + } + }); + + assertTrue(streamCompleted.await(3, TimeUnit.SECONDS), + "未收到全部终止事件时，连接关闭仍应通知流读取结束"); + } + + @Test + void shouldNotNotifyCompletionAfterDifyErrorEvent() throws Exception { + CountDownLatch errorReceived = new CountDownLatch(1); + CountDownLatch streamCompleted = new CountDownLatch(1); + DefaultDifyClient client = createClient( + "data: {\"event\":\"error\",\"status\":400,\"code\":\"invalid_param\",\"message\":\"invalid request\"}\n\n"); + + client.sendChatMessageStream(ChatMessage.builder().query("test").user("tester").build(), + new ChatflowStreamCallback() { + @Override + public void onError(ErrorEvent event) { + errorReceived.countDown(); + } + + @Override + public void onStreamComplete() { + streamCompleted.countDown(); + } + }); + + assertTrue(errorReceived.await(3, TimeUnit.SECONDS), "应分发 Dify error 事件"); + assertFalse(streamCompleted.await(300, TimeUnit.MILLISECONDS), + "Dify error 事件不应触发正常结束回调"); + } + + private DefaultDifyClient createClient(String events) { + OkHttpClient httpClient = new OkHttpClient.Builder() + .addInterceptor(chain -> new Response.Builder() + .request(chain.request()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(MediaType.parse("text/event-stream"), events)) + .build()) + .build(); + return new DefaultDifyClient("http://dify.test", "test-key", httpClient); + } + + private void assertTerminalEvents(String events, List expectedEvents) throws Exception { + List receivedEvents = Collections.synchronizedList(new ArrayList()); + CountDownLatch streamCompleted = new CountDownLatch(1); + DefaultDifyClient client = createClient(events); + + client.sendChatMessageStream(ChatMessage.builder().query("test").user("tester").build(), + new ChatflowStreamCallback() { + @Override + public void onMessageEnd(MessageEndEvent event) { + receivedEvents.add("message_end"); + } + + @Override + public void onWorkflowFinished(WorkflowFinishedEvent event) { + receivedEvents.add("workflow_finished"); + } + + @Override + public void onStreamComplete() { + receivedEvents.add("stream_complete"); + streamCompleted.countDown(); + } + }); + + assertTrue(streamCompleted.await(3, TimeUnit.SECONDS), "应在两个终止事件都收到后通知流结束"); + assertEquals(expectedEvents, receivedEvents); + } +} diff --git a/src/test/java/io/github/imfangs/dify/client/PipelineStreamCompletionTest.java b/src/test/java/io/github/imfangs/dify/client/PipelineStreamCompletionTest.java new file mode 100644 index 0000000..2f0f727 --- /dev/null +++ b/src/test/java/io/github/imfangs/dify/client/PipelineStreamCompletionTest.java @@ -0,0 +1,83 @@ +package io.github.imfangs.dify.client; + +import io.github.imfangs.dify.client.callback.WorkflowStreamCallback; +import io.github.imfangs.dify.client.event.WorkflowFinishedEvent; +import io.github.imfangs.dify.client.impl.DefaultDifyDatasetsClient; +import io.github.imfangs.dify.client.model.datasets.DatasourceNodeRunRequest; +import io.github.imfangs.dify.client.model.datasets.PipelineRunRequest; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * 验证知识库 Pipeline 流也会通知通用的流结束回调。 + */ +class PipelineStreamCompletionTest { + + private static final String WORKFLOW_FINISHED_EVENT = + "data: {\"event\":\"workflow_finished\",\"data\":{\"id\":\"workflow-1\",\"status\":\"succeeded\"}}\n\n"; + + @Test + void shouldNotifyCompletionForPipelineDatasourceNodeStream() throws Exception { + DefaultDifyDatasetsClient client = createClient(); + AtomicBoolean workflowFinishedReceived = new AtomicBoolean(); + CountDownLatch streamCompleted = new CountDownLatch(1); + + client.runPipelineDatasourceNodeStream("dataset-1", "node-1", + DatasourceNodeRunRequest.builder().build(), + createCallback(workflowFinishedReceived, streamCompleted)); + + assertTrue(streamCompleted.await(3, TimeUnit.SECONDS), "数据源节点流结束时应通知完成回调"); + assertTrue(workflowFinishedReceived.get(), "应先分发 workflow_finished 事件"); + } + + @Test + void shouldNotifyCompletionForPipelineStream() throws Exception { + DefaultDifyDatasetsClient client = createClient(); + AtomicBoolean workflowFinishedReceived = new AtomicBoolean(); + CountDownLatch streamCompleted = new CountDownLatch(1); + + client.runPipelineStream("dataset-1", PipelineRunRequest.builder().build(), + createCallback(workflowFinishedReceived, streamCompleted)); + + assertTrue(streamCompleted.await(3, TimeUnit.SECONDS), "Pipeline 流结束时应通知完成回调"); + assertTrue(workflowFinishedReceived.get(), "应先分发 workflow_finished 事件"); + } + + private WorkflowStreamCallback createCallback(AtomicBoolean workflowFinishedReceived, + CountDownLatch streamCompleted) { + return new WorkflowStreamCallback() { + @Override + public void onWorkflowFinished(WorkflowFinishedEvent event) { + workflowFinishedReceived.set(true); + } + + @Override + public void onStreamComplete() { + streamCompleted.countDown(); + } + }; + } + + private DefaultDifyDatasetsClient createClient() { + OkHttpClient httpClient = new OkHttpClient.Builder() + .addInterceptor(chain -> new Response.Builder() + .request(chain.request()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(MediaType.parse("text/event-stream"), WORKFLOW_FINISHED_EVENT)) + .build()) + .build(); + return new DefaultDifyDatasetsClient("http://dify.test", "test-key", httpClient); + } +} diff --git a/src/test/java/io/github/imfangs/dify/client/WorkflowEventsStreamCompletionTest.java b/src/test/java/io/github/imfangs/dify/client/WorkflowEventsStreamCompletionTest.java new file mode 100644 index 0000000..abb0434 --- /dev/null +++ b/src/test/java/io/github/imfangs/dify/client/WorkflowEventsStreamCompletionTest.java @@ -0,0 +1,72 @@ +package io.github.imfangs.dify.client; + +import io.github.imfangs.dify.client.callback.WorkflowStreamCallback; +import io.github.imfangs.dify.client.event.WorkflowFinishedEvent; +import io.github.imfangs.dify.client.impl.DefaultDifyClient; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * 验证工作流事件的 GET SSE 接口保留通用流结束回调。 + */ +class WorkflowEventsStreamCompletionTest { + + private static final String WORKFLOW_FINISHED_EVENT = + "data: {\"event\":\"workflow_finished\",\"data\":{\"id\":\"workflow-1\",\"status\":\"succeeded\"}}\n\n"; + + @Test + void shouldNotifyCompletionForWorkflowEventStream() throws Exception { + AtomicReference requestReference = new AtomicReference(); + DefaultDifyClient client = createClient(requestReference); + AtomicBoolean workflowFinishedReceived = new AtomicBoolean(); + CountDownLatch streamCompleted = new CountDownLatch(1); + + client.streamWorkflowEvents("workflow-1", "tester", false, false, + new WorkflowStreamCallback() { + @Override + public void onWorkflowFinished(WorkflowFinishedEvent event) { + workflowFinishedReceived.set(true); + } + + @Override + public void onStreamComplete() { + streamCompleted.countDown(); + } + }); + + assertTrue(streamCompleted.await(3, TimeUnit.SECONDS), "工作流事件流结束时应通知完成回调"); + assertTrue(workflowFinishedReceived.get(), "应先分发 workflow_finished 事件"); + assertNotNull(requestReference.get(), "应发起工作流事件请求"); + assertEquals("GET", requestReference.get().method(), "工作流事件接口必须使用 GET 请求"); + } + + private DefaultDifyClient createClient(AtomicReference requestReference) { + OkHttpClient httpClient = new OkHttpClient.Builder() + .addInterceptor(chain -> { + requestReference.set(chain.request()); + return new Response.Builder() + .request(chain.request()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(MediaType.parse("text/event-stream"), WORKFLOW_FINISHED_EVENT)) + .build(); + }) + .build(); + return new DefaultDifyClient("http://dify.test", "test-key", httpClient); + } +}