Skip to content

fix(stream): preserve Chatflow terminal events and completion #165

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Draft
SamizuHM wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(stream): preserve Chatflow terminal events and completion #165

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/main/java/io/github/imfangs/dify/client/callback/BaseStreamCallback.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,4 +33,13 @@ default void onPing(PingEvent event) {
default void onException(Throwable throwable) {
}

/**
* SDK 已完成对 SSE 响应的读取。
*
* <p>当收到客户端定义的终止事件，或服务端正常关闭 SSE 连接时触发。该回调仅表示协议流读取结束，
* 不表示 Dify 工作流或调用方业务处理成功。Dify {@code error} 事件和网络异常不会触发该回调。</p>
*/
default void onStreamComplete() {
}

}
95 changes: 90 additions & 5 deletions src/main/java/io/github/imfangs/dify/client/impl/AbstractDifyClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

/**
* 事件处理器
*/
Expand All @@ -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<Exception> 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 流式请求
*/
Expand All @@ -527,7 +563,33 @@ protected void executeGetStreamRequest(String path, LineProcessor lineProcessor,
executeStreamCall(httpRequest, lineProcessor, errorHandler);
}

/**
* 执行 GET 流式请求，并在正常结束时通知调用方。
*/
protected void executeGetStreamRequest(String path,
StreamLineProcessor lineProcessor,
Consumer<Exception> 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<Exception> errorHandler) {
executeStreamCall(httpRequest,
line -> lineProcessor.process(line) ? StreamLineResult.CONTINUE : StreamLineResult.COMPLETE,
errorHandler,
null);
}

protected void executeStreamCall(Request httpRequest,
StreamLineProcessor lineProcessor,
Consumer<Exception> errorHandler,
Runnable completionHandler) {
Call call = httpClient.newCall(httpRequest);
call.enqueue(new Callback() {
@Override
Expand Down Expand Up @@ -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) {
Expand All @@ -582,22 +654,35 @@ public void onResponse(Call call, Response response) {
* 处理一行 SSE 数据
*/
protected boolean processStreamLine(String line, BaseStreamCallback callback, Set<EventType> terminalEvents, EventProcessor eventProcessor) {
return processStreamLineWithResult(line, callback, terminalEvents, eventProcessor) == StreamLineResult.CONTINUE;
}

/**
* 处理一行 SSE 数据，并保留结束原因。
*/
protected StreamLineResult processStreamLineWithResult(String line,
BaseStreamCallback callback,
Set<EventType> 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();
try {
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);
Expand All @@ -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;
}
}
63 changes: 49 additions & 14 deletions src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -33,8 +34,9 @@
public class DefaultDifyClient extends DifyBaseClientImpl implements DifyClient {

// 流式响应相关常量
private static final Set<EventType> CHAT_TERMINAL_EVENTS = EnumSet.of(EventType.MESSAGE_END, EventType.ERROR);
private static final Set<EventType> WORKFLOW_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED, EventType.ERROR);
private static final Set<EventType> CHAT_TERMINAL_EVENTS = EnumSet.of(EventType.MESSAGE_END);
private static final Set<EventType> NO_TERMINAL_EVENTS = Collections.emptySet();
private static final Set<EventType> WORKFLOW_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED);

// API 路径常量
// 对话型应用相关路径
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -554,4 +565,28 @@ private String getFileExtension(String fileName) {
int lastDotIndex = fileName.lastIndexOf('.');
return lastDotIndex > 0 ? fileName.substring(lastDotIndex + 1) : "";
}

/**
* 跟踪 Chatflow 的两个终止事件。
*
* <p>不同 Dify 版本可能以不同顺序发送 {@code message_end} 和
* {@code workflow_finished}。只有两者都收到时，客户端才主动停止读取；若连接先关闭，
* 则由通用的流结束回调处理。</p>
*/
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;
}
}
}
12 changes: 7 additions & 5 deletions src/main/java/io/github/imfangs/dify/client/impl/DefaultDifyDatasetsClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<EventType> PIPELINE_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED, EventType.ERROR);
private static final Set<EventType> PIPELINE_TERMINAL_EVENTS = EnumSet.of(EventType.WORKFLOW_FINISHED);

@Override
public List<DatasourcePluginResponse> listPipelineDatasourcePlugins(String datasetId, Boolean isPublished) throws IOException, DifyApiException {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading