Skip to content

Commit 68dbc97

Browse files
songguocolasongguocola
committed
fix(agent/multimodal-dialog): fix safety issue and add close function
1 parent 30142dc commit 68dbc97

3 files changed

Lines changed: 228 additions & 127 deletions

File tree

‎src/main/java/com/alibaba/dashscope/multimodal/MultiModalDialog.java‎

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Queue;
2222
import java.util.UUID;
2323
import java.util.concurrent.CountDownLatch;
24+
import java.util.concurrent.TimeUnit;
2425
import java.util.concurrent.atomic.AtomicReference;
2526
import lombok.Builder;
2627
import lombok.Getter;
@@ -440,17 +441,66 @@ public void updateInfo(MultiModalRequestParam.UpdateParams updateParams) {
440441
sendTextFrame("UpdateInfo");
441442
}
442443

443-
/** Stops the MultiModalDialog. */
444+
/**
445+
* Stops the MultiModalDialog gracefully. Sends a finish message to the server and waits for the
446+
* server to acknowledge. If the server does not respond within the timeout (30s), falls back to
447+
* force close.
448+
*
449+
* <p>This method is safe to call after onError — it will not deadlock.
450+
*/
444451
public void stop() {
445-
sendFinishTaskMessage();
446-
if (stopLatch.get() != null) {
452+
// If emitter is already null (e.g., close() was called from onError), skip sending and just
453+
// ensure cleanup
454+
boolean sent;
455+
synchronized (MultiModalDialog.this) {
456+
sent = (conversationEmitter != null);
457+
}
458+
if (sent) {
459+
sendFinishTaskMessage();
460+
}
461+
462+
CountDownLatch latch = stopLatch.get();
463+
if (latch != null) {
447464
try {
448-
stopLatch.get().await();
465+
// Use timeout to prevent deadlock: if server doesn't respond in 30s, force close
466+
boolean completed = latch.await(30, TimeUnit.SECONDS);
467+
if (!completed) {
468+
log.warn("stop() timed out waiting for server acknowledgement, forcing close");
469+
close();
470+
}
449471
} catch (InterruptedException ignored) {
472+
close();
450473
}
451474
}
452475
}
453476

477+
/**
478+
* Force closes the dialog immediately without sending any message to the server. Safe to call
479+
* from any callback thread (including onError/onStopped). Does not block.
480+
*
481+
* <p>This method: - Nullifies emitter to prevent further data sending - Disposes upstream
482+
* subscription to stop sendBinaryWithRetry - Forces WebSocket cancel (non-blocking) - Sets
483+
* isClosed to short-circuit any reconnection loops in progress - Releases stopLatch to unblock
484+
* any thread waiting in stop()
485+
*/
486+
public void close() {
487+
// Nullify emitter to prevent further data sending
488+
synchronized (MultiModalDialog.this) {
489+
conversationEmitter = null;
490+
}
491+
// Force cancel: disposes upstream subscription + cancels WebSocket (non-blocking)
492+
// Also sets isClosed=true which breaks reconnection loops in
493+
// establishWebSocketClient/sendTextWithRetry/sendBinaryWithRetry
494+
duplexApi.cancel();
495+
// Release stopLatch to prevent stop() from blocking forever.
496+
// This is critical: if close() is called from onError callback, any thread waiting
497+
// in stop() will be unblocked immediately instead of deadlocking.
498+
CountDownLatch latch = stopLatch.get();
499+
if (latch != null) {
500+
latch.countDown();
501+
}
502+
}
503+
454504
/**
455505
* Gets current dialogue state.
456506
*

‎src/main/java/com/alibaba/dashscope/protocol/okhttp/OkHttpWebSocketClient.java‎

Lines changed: 110 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import io.reactivex.Flowable;
1818
import io.reactivex.FlowableEmitter;
1919
import io.reactivex.Observable;
20+
import io.reactivex.disposables.Disposable;
2021
import io.reactivex.functions.Action;
2122
import java.io.IOException;
2223
import java.nio.ByteBuffer;
@@ -51,6 +52,9 @@ public class OkHttpWebSocketClient extends WebSocketListener
5152

5253
private AtomicBoolean passTaskStarted = new AtomicBoolean(false);
5354

55+
// Disposable for the streaming data subscription, used to cancel upstream when stopping
56+
protected volatile Disposable streamingDataDisposable;
57+
5458
public OkHttpWebSocketClient(OkHttpClient client, boolean passTaskStarted) {
5559
this.client = client;
5660
this.passTaskStarted.set(passTaskStarted);
@@ -97,6 +101,13 @@ public boolean close(int code, String reason) {
97101
}
98102

99103
public void cancel() {
104+
// Set isClosed BEFORE cancel to suppress onFailure error propagation
105+
isClosed.set(true);
106+
// Dispose upstream subscription to stop sending data
107+
Disposable d = streamingDataDisposable;
108+
if (d != null && !d.isDisposed()) {
109+
d.dispose();
110+
}
100111
if (webSocketClient != null) {
101112
webSocketClient.cancel();
102113
}
@@ -111,6 +122,11 @@ private void establishWebSocketClient(
111122
int reconnectionTimes = 0;
112123
String errorMessage = "";
113124
while (reconnectionTimes < MAX_CONNECTION_TIMES) {
125+
// Bail out immediately if cancel() has been called
126+
if (isClosed.get()) {
127+
log.debug("Connection cancelled, stop reconnecting.");
128+
return;
129+
}
114130
try {
115131
Flowable<DashScopeResult> flowable =
116132
Flowable.<DashScopeResult>create(
@@ -144,9 +160,17 @@ private void establishWebSocketClient(
144160
} else if (errorMessage.contains(Constants.NO_API_KEY_ERROR)) {
145161
throw ex;
146162
}
163+
// Check again before sleeping
164+
if (isClosed.get()) {
165+
log.debug("Connection cancelled during retry, stop reconnecting.");
166+
return;
167+
}
147168
try {
148169
Thread.sleep(10000);
149-
} catch (InterruptedException e) {;
170+
} catch (InterruptedException e) {
171+
// Respect interruption - exit the loop
172+
Thread.currentThread().interrupt();
173+
return;
150174
}
151175
}
152176
}
@@ -379,10 +403,18 @@ protected void sendTextWithRetry(
379403
String workspace,
380404
Map<String, String> customHeaders,
381405
String baseWebSocketUrl) {
406+
// Guard: skip if already cancelled
407+
if (isClosed.get()) {
408+
log.debug("sendTextWithRetry skipped: connection already closed.");
409+
return;
410+
}
382411
// simple retry with fixed delay, no strategy
383412
if (!isOpen.get()) {
384413
establishWebSocketClient(apiKey, isSecurityCheck, workspace, customHeaders, baseWebSocketUrl);
385414
}
415+
if (isClosed.get()) {
416+
return;
417+
}
386418
int maxRetries = 3;
387419
if (passTaskStarted.get()) {
388420
// when pass througn task started, no need to retry.
@@ -395,6 +427,9 @@ protected void sendTextWithRetry(
395427
}
396428
int retryCount = 0;
397429
while (retryCount < maxRetries) {
430+
if (isClosed.get()) {
431+
return;
432+
}
398433
log.debug("Sending message: " + message);
399434
Boolean isOk = webSocketClient.send(message);
400435
if (isOk) {
@@ -418,12 +453,22 @@ protected void sendBinaryWithRetry(
418453
String workspace,
419454
Map<String, String> customHeaders,
420455
String baseWebSocketUrl) {
456+
// Guard: skip if already cancelled
457+
if (isClosed.get()) {
458+
return;
459+
}
421460
if (!isOpen.get()) {
422461
establishWebSocketClient(apiKey, isSecurityCheck, workspace, customHeaders, baseWebSocketUrl);
423462
}
463+
if (isClosed.get()) {
464+
return;
465+
}
424466
int maxRetries = 3;
425467
int retryCount = 0;
426468
while (retryCount < maxRetries) {
469+
if (isClosed.get()) {
470+
return;
471+
}
427472
Boolean isOk = webSocketClient.send(message);
428473
if (isOk) {
429474
break;
@@ -584,67 +629,70 @@ protected CompletableFuture<Void> sendStreamRequest(FullDuplexRequest req) {
584629
req.getBaseWebSocketUrl());
585630

586631
Flowable<Object> streamingData = req.getStreamingData();
587-
streamingData.subscribe(
588-
data -> {
589-
try {
590-
if (data instanceof String) {
591-
JsonObject continueData = req.getContinueMessage((String) data, taskId);
592-
sendTextWithRetry(
593-
req.getApiKey(),
594-
req.isSecurityCheck(),
595-
JsonUtils.toJson(continueData),
596-
req.getWorkspace(),
597-
req.getHeaders(),
598-
req.getBaseWebSocketUrl());
599-
} else if (data instanceof byte[]) {
600-
sendBinaryWithRetry(
601-
req.getApiKey(),
602-
req.isSecurityCheck(),
603-
ByteString.of((byte[]) data),
604-
req.getWorkspace(),
605-
req.getHeaders(),
606-
req.getBaseWebSocketUrl());
607-
} else if (data instanceof ByteBuffer) {
608-
sendBinaryWithRetry(
609-
req.getApiKey(),
610-
req.isSecurityCheck(),
611-
ByteString.of((ByteBuffer) data),
612-
req.getWorkspace(),
613-
req.getHeaders(),
614-
req.getBaseWebSocketUrl());
615-
} else {
616-
JsonObject continueData = req.getContinueMessage(data, taskId);
617-
sendTextWithRetry(
618-
req.getApiKey(),
619-
req.isSecurityCheck(),
620-
JsonUtils.toJson(continueData),
621-
req.getWorkspace(),
622-
req.getHeaders(),
623-
req.getBaseWebSocketUrl());
624-
}
625-
} catch (Throwable ex) {
626-
log.error(
627-
StringUtils.format("sendStreamData exception: %s", ex.getMessage()));
628-
responseEmitter.onError(ex);
629-
}
630-
},
631-
err -> {
632-
log.error(StringUtils.format("Get stream data error!"));
633-
responseEmitter.onError(err);
634-
},
635-
new Action() {
636-
@Override
637-
public void run() throws Exception {
638-
log.debug(StringUtils.format("Stream data send completed!"));
639-
sendTextWithRetry(
640-
req.getApiKey(),
641-
req.isSecurityCheck(),
642-
JsonUtils.toJson(req.getFinishedTaskMessage(taskId)),
643-
req.getWorkspace(),
644-
req.getHeaders(),
645-
req.getBaseWebSocketUrl());
646-
}
647-
});
632+
streamingDataDisposable =
633+
streamingData.subscribe(
634+
data -> {
635+
try {
636+
if (data instanceof String) {
637+
JsonObject continueData =
638+
req.getContinueMessage((String) data, taskId);
639+
sendTextWithRetry(
640+
req.getApiKey(),
641+
req.isSecurityCheck(),
642+
JsonUtils.toJson(continueData),
643+
req.getWorkspace(),
644+
req.getHeaders(),
645+
req.getBaseWebSocketUrl());
646+
} else if (data instanceof byte[]) {
647+
sendBinaryWithRetry(
648+
req.getApiKey(),
649+
req.isSecurityCheck(),
650+
ByteString.of((byte[]) data),
651+
req.getWorkspace(),
652+
req.getHeaders(),
653+
req.getBaseWebSocketUrl());
654+
} else if (data instanceof ByteBuffer) {
655+
sendBinaryWithRetry(
656+
req.getApiKey(),
657+
req.isSecurityCheck(),
658+
ByteString.of((ByteBuffer) data),
659+
req.getWorkspace(),
660+
req.getHeaders(),
661+
req.getBaseWebSocketUrl());
662+
} else {
663+
JsonObject continueData = req.getContinueMessage(data, taskId);
664+
sendTextWithRetry(
665+
req.getApiKey(),
666+
req.isSecurityCheck(),
667+
JsonUtils.toJson(continueData),
668+
req.getWorkspace(),
669+
req.getHeaders(),
670+
req.getBaseWebSocketUrl());
671+
}
672+
} catch (Throwable ex) {
673+
log.error(
674+
StringUtils.format(
675+
"sendStreamData exception: %s", ex.getMessage()));
676+
responseEmitter.onError(ex);
677+
}
678+
},
679+
err -> {
680+
log.error(StringUtils.format("Get stream data error!"));
681+
responseEmitter.onError(err);
682+
},
683+
new Action() {
684+
@Override
685+
public void run() throws Exception {
686+
log.debug(StringUtils.format("Stream data send completed!"));
687+
sendTextWithRetry(
688+
req.getApiKey(),
689+
req.isSecurityCheck(),
690+
JsonUtils.toJson(req.getFinishedTaskMessage(taskId)),
691+
req.getWorkspace(),
692+
req.getHeaders(),
693+
req.getBaseWebSocketUrl());
694+
}
695+
});
648696
} catch (Throwable ex) {
649697
log.error(StringUtils.format("sendStreamData exception: %s", ex.getMessage()));
650698
responseEmitter.onError(ex);

0 commit comments

Comments
 (0)