# WebSocket 任务通知接收示例

## 一、WebSocket 配置说明

### 后端配置
- **连接端点**: `/user/websocket`（支持 SockJS 备用方案）
- **用户前缀**: `/user`
- **订阅目的地**: `/user/queue/tasks-progress`
- **协议**: STOMP over WebSocket

### 消息格式（TaskProgressDto）
```typescript
interface TaskProgressDto {
  taskNo: string;           // 任务编号
  status: string;           // 任务状态: created/queued/processing/completed/failed
  progress: number;         // 进度百分比 0-100
  message: string;          // 进度消息
  resultUrl?: string;       // 结果URL（完成时）
  errorMessage?: string;    // 错误信息（失败时）
}
```

---

## 二、前端依赖安装

### 使用 npm
```bash
npm install @stomp/stompjs sockjs-client
```

### 使用 yarn
```bash
yarn add @stomp/stompjs sockjs-client
```

### CDN 引入（HTML）
```html
<script src="https://cdn.jsdelivr.net/npm/sockjs-client@1.6.1/dist/sockjs.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@stomp/stompjs@7.0.0/bundles/stomp.umd.min.js"></script>
```

---

## 三、基础连接示例

### 1. 原生 JavaScript + STOMP.js

```javascript
// 引入依赖（如果使用模块化）
import SockJS from 'sockjs-client';
import { Client } from '@stomp/stompjs';

// WebSocket 配置
const WEBSOCKET_URL = 'http://localhost:8081/ws';
const AUTH_TOKEN = 'YOUR_JWT_TOKEN';

// 创建 STOMP 客户端
const stompClient = new Client({
  // 使用 SockJS 作为 WebSocket 实现
  webSocketFactory: () => new SockJS(WEBSOCKET_URL),
  
  // 连接头（携带认证信息）
  connectHeaders: {
    Authorization: `Bearer ${AUTH_TOKEN}`
  },
  
  // 连接成功回调
  onConnect: (frame) => {
    console.log('WebSocket 连接成功:', frame);
    
    // 订阅任务进度更新
    stompClient.subscribe('/user/queue/tasks-progress', (message) => {
      const notification = JSON.parse(message.body);
      console.log('收到任务通知:', notification);
      
      // 处理通知
      handleTaskNotification(notification);
    });
  },
  
  // 连接失败回调
  onStompError: (frame) => {
    console.error('STOMP 错误:', frame);
  },
  
  // WebSocket 错误回调
  onWebSocketError: (error) => {
    console.error('WebSocket 错误:', error);
  },
  
  // WebSocket 关闭回调
  onWebSocketClose: (event) => {
    console.log('WebSocket 连接关闭:', event);
  },
  
  // 自动重连配置
  reconnectDelay: 5000,  // 5秒后重连
  heartbeatIncoming: 4000,
  heartbeatOutgoing: 4000
});

// 激活连接
stompClient.activate();

// 处理任务通知
function handleTaskNotification(notification) {
  const { taskNo, status, progress, message, resultUrl, errorMessage } = notification;
  
  switch(status) {
    case 'created':
      console.log(`[${taskNo}] 任务已创建`);
      break;
    case 'queued':
      console.log(`[${taskNo}] 任务排队中: ${message}`);
      break;
    case 'processing':
      console.log(`[${taskNo}] 处理中 ${progress}%: ${message}`);
      updateProgressBar(taskNo, progress);
      break;
    case 'completed':
      console.log(`[${taskNo}] 任务完成: ${resultUrl}`);
      showCompletedNotification(taskNo, resultUrl);
      break;
    case 'failed':
      console.error(`[${taskNo}] 任务失败: ${errorMessage || message}`);
      showErrorNotification(taskNo, errorMessage || message);
      break;
  }
}

// 断开连接
function disconnect() {
  if (stompClient.connected) {
    stompClient.deactivate();
    console.log('WebSocket 已断开');
  }
}
```

---

## 四、完整封装类（推荐）

### TaskWebSocketClient.js

```javascript
import SockJS from 'sockjs-client';
import { Client } from '@stomp/stompjs';

class TaskWebSocketClient {
  constructor(baseUrl = 'http://localhost:8081', token) {
    this.baseUrl = baseUrl;
    this.token = token;
    this.client = null;
    this.listeners = {
      onTaskUpdate: [],
      onConnect: [],
      onDisconnect: [],
      onError: []
    };
  }
  
  /**
   * 连接 WebSocket
   */
  connect() {
    if (this.client && this.client.connected) {
      console.warn('WebSocket 已连接');
      return Promise.resolve();
    }
    
    return new Promise((resolve, reject) => {
      this.client = new Client({
        webSocketFactory: () => new SockJS(`${this.baseUrl}/ws`),
        
        connectHeaders: {
          Authorization: `Bearer ${this.token}`
        },
        
        onConnect: (frame) => {
          console.log('✅ WebSocket 连接成功');
          
          // 订阅任务进度更新
          this.client.subscribe('/user/queue/tasks-progress', (message) => {
            try {
              const notification = JSON.parse(message.body);
              this._notifyListeners('onTaskUpdate', notification);
            } catch (error) {
              console.error('解析消息失败:', error);
            }
          });
          
          this._notifyListeners('onConnect', frame);
          resolve(frame);
        },
        
        onStompError: (frame) => {
          console.error('❌ STOMP 错误:', frame);
          this._notifyListeners('onError', frame);
          reject(frame);
        },
        
        onWebSocketError: (error) => {
          console.error('❌ WebSocket 错误:', error);
          this._notifyListeners('onError', error);
        },
        
        onWebSocketClose: (event) => {
          console.log('🔌 WebSocket 连接关闭');
          this._notifyListeners('onDisconnect', event);
        },
        
        reconnectDelay: 5000,
        heartbeatIncoming: 4000,
        heartbeatOutgoing: 4000
      });
      
      this.client.activate();
    });
  }
  
  /**
   * 断开 WebSocket
   */
  disconnect() {
    if (this.client && this.client.connected) {
      this.client.deactivate();
      console.log('WebSocket 已断开');
    }
  }
  
  /**
   * 添加任务更新监听器
   */
  onTaskUpdate(callback) {
    this.listeners.onTaskUpdate.push(callback);
    return () => this._removeListener('onTaskUpdate', callback);
  }
  
  /**
   * 添加连接监听器
   */
  onConnect(callback) {
    this.listeners.onConnect.push(callback);
    return () => this._removeListener('onConnect', callback);
  }
  
  /**
   * 添加断开连接监听器
   */
  onDisconnect(callback) {
    this.listeners.onDisconnect.push(callback);
    return () => this._removeListener('onDisconnect', callback);
  }
  
  /**
   * 添加错误监听器
   */
  onError(callback) {
    this.listeners.onError.push(callback);
    return () => this._removeListener('onError', callback);
  }
  
  /**
   * 通知所有监听器
   */
  _notifyListeners(event, data) {
    this.listeners[event].forEach(callback => {
      try {
        callback(data);
      } catch (error) {
        console.error(`监听器执行错误 (${event}):`, error);
      }
    });
  }
  
  /**
   * 移除监听器
   */
  _removeListener(event, callback) {
    const index = this.listeners[event].indexOf(callback);
    if (index > -1) {
      this.listeners[event].splice(index, 1);
    }
  }
  
  /**
   * 检查连接状态
   */
  isConnected() {
    return this.client && this.client.connected;
  }
}

export default TaskWebSocketClient;
```

---

## 五、使用示例

### 1. React 组件示例

```jsx
import React, { useEffect, useState } from 'react';
import TaskWebSocketClient from './TaskWebSocketClient';

function TaskMonitor() {
  const [tasks, setTasks] = useState({});
  const [wsClient, setWsClient] = useState(null);
  
  useEffect(() => {
    // 获取 Token（从 localStorage 或其他地方）
    const token = localStorage.getItem('jwt_token');
    
    // 创建 WebSocket 客户端
    const client = new TaskWebSocketClient('http://localhost:8081', token);
    
    // 监听任务更新
    const unsubscribe = client.onTaskUpdate((notification) => {
      console.log('任务更新:', notification);
      
      // 更新任务状态
      setTasks(prev => ({
        ...prev,
        [notification.taskNo]: notification
      }));
      
      // 根据状态显示不同提示
      if (notification.status === 'completed') {
        showSuccessToast(`任务 ${notification.taskNo} 已完成！`);
      } else if (notification.status === 'failed') {
        showErrorToast(`任务 ${notification.taskNo} 失败: ${notification.errorMessage}`);
      }
    });
    
    // 连接 WebSocket
    client.connect().catch(error => {
      console.error('连接失败:', error);
    });
    
    setWsClient(client);
    
    // 清理函数
    return () => {
      unsubscribe();
      client.disconnect();
    };
  }, []);
  
  return (
    <div>
      <h2>任务监控</h2>
      {Object.values(tasks).map(task => (
        <div key={task.taskNo} className="task-card">
          <h3>{task.taskNo}</h3>
          <p>状态: {task.status}</p>
          <p>进度: {task.progress}%</p>
          <p>{task.message}</p>
          {task.resultUrl && (
            <a href={task.resultUrl} target="_blank" rel="noopener noreferrer">
              查看结果
            </a>
          )}
        </div>
      ))}
    </div>
  );
}

export default TaskMonitor;
```

### 2. Vue 3 组件示例

```vue
<template>
  <div class="task-monitor">
    <h2>任务监控</h2>
    <div v-for="task in taskList" :key="task.taskNo" class="task-card">
      <h3>{{ task.taskNo }}</h3>
      <p>状态: {{ task.status }}</p>
      <div v-if="task.status === 'processing'">
        <progress :value="task.progress" max="100"></progress>
        <span>{{ task.progress }}%</span>
      </div>
      <p>{{ task.message }}</p>
      <a v-if="task.resultUrl" :href="task.resultUrl" target="_blank">
        查看结果
      </a>
    </div>
  </div>
</template>

<script setup>
import { ref, onMounted, onUnmounted } from 'vue';
import TaskWebSocketClient from './TaskWebSocketClient';

const tasks = ref({});
const taskList = computed(() => Object.values(tasks.value));
let wsClient = null;

onMounted(() => {
  const token = localStorage.getItem('jwt_token');
  wsClient = new TaskWebSocketClient('http://localhost:8081', token);
  
  // 监听任务更新
  wsClient.onTaskUpdate((notification) => {
    tasks.value[notification.taskNo] = notification;
    
    if (notification.status === 'completed') {
      ElMessage.success(`任务 ${notification.taskNo} 已完成！`);
    } else if (notification.status === 'failed') {
      ElMessage.error(`任务失败: ${notification.errorMessage}`);
    }
  });
  
  // 连接
  wsClient.connect();
});

onUnmounted(() => {
  if (wsClient) {
    wsClient.disconnect();
  }
});
</script>
```

### 3. 原生 JavaScript 完整示例

```html
<!DOCTYPE html>
<html>
<head>
  <title>任务监控</title>
  <script src="https://cdn.jsdelivr.net/npm/sockjs-client@1.6.1/dist/sockjs.min.js"></script>
  <script src="https://cdn.jsdelivr.net/npm/@stomp/stompjs@7.0.0/bundles/stomp.umd.min.js"></script>
</head>
<body>
  <h1>AI 任务实时监控</h1>
  <div id="connection-status">未连接</div>
  <div id="tasks-container"></div>
  
  <script>
    const AUTH_TOKEN = localStorage.getItem('jwt_token');
    const tasks = {};
    
    // 创建 STOMP 客户端
    const client = new StompJs.Client({
      webSocketFactory: () => new SockJS('http://localhost:8081/ws'),
      connectHeaders: {
        Authorization: `Bearer ${AUTH_TOKEN}`
      },
      onConnect: (frame) => {
        document.getElementById('connection-status').textContent = '✅ 已连接';
        
        // 订阅任务通知
        client.subscribe('/user/queue/tasks-progress', (message) => {
          const notification = JSON.parse(message.body);
          handleTaskNotification(notification);
        });
      },
      onStompError: (frame) => {
        document.getElementById('connection-status').textContent = '❌ 连接错误';
        console.error('STOMP 错误:', frame);
      },
      reconnectDelay: 5000,
      heartbeatIncoming: 4000,
      heartbeatOutgoing: 4000
    });
    
    // 处理任务通知
    function handleTaskNotification(notification) {
      const { taskNo, status, progress, message, resultUrl, errorMessage } = notification;
      
      // 更新任务数据
      tasks[taskNo] = notification;
      
      // 渲染任务列表
      renderTasks();
      
      // 显示浏览器通知（如果支持）
      if (status === 'completed' && 'Notification' in window && Notification.permission === 'granted') {
        new Notification('任务完成', {
          body: `任务 ${taskNo} 已成功完成！`,
          icon: '/icon.png'
        });
      }
    }
    
    // 渲染任务列表
    function renderTasks() {
      const container = document.getElementById('tasks-container');
      container.innerHTML = '';
      
      Object.values(tasks).forEach(task => {
        const taskDiv = document.createElement('div');
        taskDiv.className = 'task-card';
        taskDiv.innerHTML = `
          <h3>${task.taskNo}</h3>
          <p>状态: ${task.status}</p>
          <p>进度: ${task.progress}%</p>
          <progress value="${task.progress}" max="100"></progress>
          <p>${task.message}</p>
          ${task.resultUrl ? `<a href="${task.resultUrl}" target="_blank">查看结果</a>` : ''}
          ${task.errorMessage ? `<p style="color:red">错误: ${task.errorMessage}</p>` : ''}
        `;
        container.appendChild(taskDiv);
      });
    }
    
    // 激活连接
    client.activate();
    
    // 页面关闭时断开连接
    window.addEventListener('beforeunload', () => {
      if (client.connected) {
        client.deactivate();
      }
    });
    
    // 请求浏览器通知权限
    if ('Notification' in window && Notification.permission === 'default') {
      Notification.requestPermission();
    }
  </script>
  
  <style>
    .task-card {
      border: 1px solid #ddd;
      padding: 15px;
      margin: 10px 0;
      border-radius: 8px;
    }
    progress {
      width: 100%;
      height: 20px;
    }
  </style>
</body>
</html>
```

---

## 六、完整业务流程示例

```javascript
import TaskWebSocketClient from './TaskWebSocketClient';
import SuChuangImageGenerator from './SuChuangImageGenerator';

class AITaskManager {
  constructor(token, baseUrl = 'http://localhost:8081') {
    this.wsClient = new TaskWebSocketClient(baseUrl, token);
    this.generator = new SuChuangImageGenerator(token, baseUrl);
    this.pendingTasks = new Map();
  }
  
  /**
   * 初始化（连接 WebSocket 并设置监听器）
   */
  async init() {
    // 监听任务更新
    this.wsClient.onTaskUpdate((notification) => {
      this.handleTaskUpdate(notification);
    });
    
    // 连接 WebSocket
    await this.wsClient.connect();
    console.log('AI 任务管理器已初始化');
  }
  
  /**
   * 提交文生图任务
   */
  async submitTextToImage(prompt, aspectRatio = '1:1', onProgress, onComplete, onError) {
    try {
      // 提交任务
      const taskNo = await this.generator.generateImage(prompt, aspectRatio);
      
      // 注册回调
      this.pendingTasks.set(taskNo, { onProgress, onComplete, onError });
      
      return taskNo;
    } catch (error) {
      if (onError) onError(error);
      throw error;
    }
  }
  
  /**
   * 提交图生图任务
   */
  async submitImageToImage(prompt, imageUrl, aspectRatio = '1:1', onProgress, onComplete, onError) {
    try {
      const taskNo = await this.generator.transformImage(prompt, imageUrl, aspectRatio);
      this.pendingTasks.set(taskNo, { onProgress, onComplete, onError });
      return taskNo;
    } catch (error) {
      if (onError) onError(error);
      throw error;
    }
  }
  
  /**
   * 处理任务更新
   */
  handleTaskUpdate(notification) {
    const { taskNo, status, progress, message, resultUrl, errorMessage } = notification;
    const callbacks = this.pendingTasks.get(taskNo);
    
    if (!callbacks) return;
    
    const { onProgress, onComplete, onError } = callbacks;
    
    switch(status) {
      case 'processing':
        if (onProgress) onProgress(progress, message);
        break;
      
      case 'completed':
        if (onComplete) onComplete(resultUrl);
        this.pendingTasks.delete(taskNo);
        break;
      
      case 'failed':
        if (onError) onError(new Error(errorMessage || message));
        this.pendingTasks.delete(taskNo);
        break;
    }
  }
  
  /**
   * 清理资源
   */
  destroy() {
    this.wsClient.disconnect();
    this.pendingTasks.clear();
  }
}

// 使用示例
const taskManager = new AITaskManager(YOUR_JWT_TOKEN);

// 初始化
await taskManager.init();

// 提交任务并监听进度
const taskNo = await taskManager.submitTextToImage(
  '一只可爱的柴犬',
  '1:1',
  (progress, message) => {
    console.log(`进度: ${progress}% - ${message}`);
    updateProgressBar(progress);
  },
  (resultUrl) => {
    console.log('生成完成:', resultUrl);
    displayImage(resultUrl);
  },
  (error) => {
    console.error('生成失败:', error);
    showErrorMessage(error.message);
  }
);

console.log('任务已提交:', taskNo);
```

---

## 七、常见问题

### 1. 跨域问题
如果前端和后端不在同一域名，确保后端已配置 CORS：

```java
// WebSocketConfig.java 中已配置
registry.addEndpoint("/ws").setAllowedOriginPatterns("*").withSockJS();
```

### 2. 认证失败
确保在连接时传递了正确的 JWT Token：

```javascript
connectHeaders: {
  Authorization: `Bearer ${YOUR_JWT_TOKEN}`
}
```

### 3. 连接断开自动重连
STOMP 客户端已配置自动重连：

```javascript
reconnectDelay: 5000  // 5秒后自动重连
```

### 4. 心跳检测
防止连接超时：

```javascript
heartbeatIncoming: 4000,  // 接收心跳间隔
heartbeatOutgoing: 4000   // 发送心跳间隔
```

---

## 八、调试技巧

### 1. 开启详细日志
```javascript
const client = new Client({
  // ... 其他配置
  debug: (str) => {
    console.log('STOMP Debug:', str);
  }
});
```

### 2. 监控连接状态
```javascript
client.onConnect = () => console.log('✅ 已连接');
client.onDisconnect = () => console.log('🔌 已断开');
client.onWebSocketError = (error) => console.error('❌ 错误:', error);
```

### 3. 测试消息接收
```javascript
client.subscribe('/user/queue/tasks-progress', (message) => {
  console.log('收到原始消息:', message.body);
  const data = JSON.parse(message.body);
  console.log('解析后的数据:', data);
});
```

---

## 九、安全建议

1. **不要在前端暴露敏感信息**
   - Token 应通过 localStorage 或 sessionStorage 安全存储
   - 避免在 URL 中传递 Token

2. **设置合理的超时时间**
   - 避免长时间保持空闲连接

3. **处理连接断开**
   - 实现重连逻辑
   - 提示用户连接状态

4. **验证消息来源**
   - 确认 taskNo 是否为当前用户的任务

---

## 十、完整目录结构

```
src/
├── websocket/
│   ├── TaskWebSocketClient.js       # WebSocket 封装类
│   └── AITaskManager.js             # 任务管理器
├── api/
│   └── SuChuangImageGenerator.js    # 任务提交API
├── components/
│   ├── TaskMonitor.jsx              # React 任务监控组件
│   └── TaskMonitor.vue              # Vue 任务监控组件
└── utils/
    └── notification.js              # 浏览器通知工具
```

---

## 总结

前端通过以下步骤接收 WebSocket 通知：

1. **连接**: `new SockJS('http://localhost:8081/ws')`
2. **认证**: 在 `connectHeaders` 中传递 JWT Token
3. **订阅**: `client.subscribe('/user/queue/tasks-progress', callback)`
4. **处理**: 根据 `status` 字段处理不同状态的通知

**关键点**：
- ✅ 端点以 `/ws` 开头（不是 `/user/ws`）
- ✅ 订阅地址为 `/user/queue/tasks-progress`
- ✅ Spring 会自动将消息路由到当前用户
- ✅ 支持自动重连和心跳检测