001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.model.client.impl;
017
018import com.agentsflex.core.model.chat.ChatConfig;
019import com.agentsflex.core.model.chat.log.ChatMessageLogger;
020import com.agentsflex.core.model.client.OkHttpClientUtil;
021import com.agentsflex.core.model.client.StreamClient;
022import com.agentsflex.core.model.client.StreamClientListener;
023import okhttp3.*;
024import okio.ByteString;
025import org.jetbrains.annotations.NotNull;
026import org.slf4j.Logger;
027import org.slf4j.LoggerFactory;
028
029import java.util.Map;
030
031public class WebSocketClient extends WebSocketListener implements StreamClient {
032    private static final Logger log = LoggerFactory.getLogger(WebSocketClient.class);
033    private OkHttpClient okHttpClient;
034    private WebSocket webSocket;
035    private StreamClientListener listener;
036    private ChatConfig config;
037    private boolean isStop = false;
038    private String payload;
039
040
041    public WebSocketClient() {
042        this(OkHttpClientUtil.buildDefaultClient());
043    }
044
045    public WebSocketClient(OkHttpClient okHttpClient) {
046        if (okHttpClient == null) {
047            throw new IllegalArgumentException("OkHttpClient must not be null");
048        }
049        this.okHttpClient = okHttpClient;
050    }
051
052    public OkHttpClient getOkHttpClient() {
053        return okHttpClient;
054    }
055
056    public void setOkHttpClient(OkHttpClient okHttpClient) {
057        this.okHttpClient = okHttpClient;
058    }
059
060    @Override
061    public void start(String url, Map<String, String> headers, String payload, StreamClientListener listener, ChatConfig config) {
062        if (isStop) {
063            throw new IllegalStateException("WebSocketClient has been stopped and cannot be reused.");
064        }
065
066        this.listener = listener;
067        this.payload = payload;
068        this.config = config;
069        this.isStop = false;
070
071        Request.Builder builder = new Request.Builder().url(url);
072        if (headers != null && !headers.isEmpty()) {
073            headers.forEach(builder::addHeader);
074        }
075
076        Request request = builder.build();
077
078        // 创建 WebSocket 连接
079        this.webSocket = okHttpClient.newWebSocket(request, this);
080        ChatMessageLogger.logRequest(config, payload);
081    }
082
083    @Override
084    public void stop() {
085        closeWebSocketAndNotify();
086    }
087
088
089    //webSocket events
090    @Override
091    public void onOpen(WebSocket webSocket, Response response) {
092        webSocket.send(payload);
093        this.listener.onStart(this);
094    }
095
096    @Override
097    public void onMessage(WebSocket webSocket, String text) {
098        ChatMessageLogger.logResponse(config, text);
099        this.listener.onMessage(this, text);
100    }
101
102    @Override
103    public void onMessage(WebSocket webSocket, ByteString bytes) {
104        this.onMessage(webSocket, bytes.utf8());
105    }
106
107    @Override
108    public void onClosing(WebSocket webSocket, int code, String reason) {
109        closeWebSocketAndNotify();
110    }
111
112    @Override
113    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
114        try {
115            Throwable failureThrowable = Util.getFailureThrowable(t, response);
116            this.listener.onFailure(this, failureThrowable);
117        } finally {
118            closeWebSocketAndNotify();
119        }
120    }
121
122    @Override
123    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
124        closeWebSocketAndNotify();
125    }
126
127
128    /**
129     * 安全关闭 WebSocket 并通知监听器（确保只执行一次）
130     */
131    private void closeWebSocketAndNotify() {
132        if (isStop) return;
133        synchronized (this) {
134            if (isStop) return;
135            isStop = true;
136
137            // 先通知 onStop
138            if (this.listener != null) {
139                try {
140                    this.listener.onStop(this);
141                } catch (Exception e) {
142                    log.warn(e.getMessage(), e);
143                }
144            }
145
146            // 再关闭 WebSocket（幂等：close 多次无害，但避免空指针）
147            if (this.webSocket != null) {
148                try {
149                    this.webSocket.close(1000, ""); // 正常关闭
150                } catch (Exception e) {
151                    // 忽略关闭异常（连接可能已断）
152                } finally {
153                    this.webSocket = null;
154                }
155            }
156        }
157    }
158}