| 1 | package bot |
| 2 | |
| 3 | import ( |
| 4 | "crypto/sha256" |
| 5 | "encoding/hex" |
| 6 | "fmt" |
| 7 | "strings" |
| 8 | "sync" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | const ( |
| 13 | QueueModeSteer = "steer" |
| 14 | QueueModeFollowup = "followup" |
| 15 | QueueModeCollect = "collect" |
| 16 | QueueModeInterrupt = "interrupt" |
| 17 | |
| 18 | QueueDropSummarize = "summarize" |
| 19 | QueueDropOld = "old" |
| 20 | QueueDropNew = "new" |
| 21 | |
| 22 | DefaultQueueCap = 20 |
| 23 | ) |
| 24 | |
| 25 | type QueueOptions struct { |
| 26 | Mode string |
| 27 | Cap int |
| 28 | Drop string |
| 29 | } |
| 30 | |
| 31 | type QueueResult struct { |
| 32 | Acquired bool |
| 33 | Queued bool |
| 34 | Rejected bool |
| 35 | Dropped bool |
| 36 | Pending int |
| 37 | Mode string |
| 38 | } |
| 39 | |
| 40 | type QueueSnapshot struct { |
| 41 | Active int |
| 42 | Pending int |
| 43 | Dropped int |
| 44 | Sessions int |
| 45 | } |
| 46 | |
| 47 | func NormalizeQueueMode(mode string) string { |
| 48 | switch strings.ToLower(strings.TrimSpace(mode)) { |
| 49 | case QueueModeSteer: |
| 50 | return QueueModeSteer |
| 51 | case QueueModeFollowup: |
| 52 | return QueueModeFollowup |
| 53 | case QueueModeCollect: |
| 54 | return QueueModeCollect |
| 55 | case QueueModeInterrupt: |
| 56 | return QueueModeInterrupt |
| 57 | default: |
| 58 | return QueueModeSteer |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | func NormalizeOptionalQueueMode(mode string) string { |
| 63 | switch strings.ToLower(strings.TrimSpace(mode)) { |
| 64 | case QueueModeSteer: |
| 65 | return QueueModeSteer |
| 66 | case QueueModeFollowup: |
| 67 | return QueueModeFollowup |
| 68 | case QueueModeCollect: |
| 69 | return QueueModeCollect |
| 70 | case QueueModeInterrupt: |
| 71 | return QueueModeInterrupt |
| 72 | default: |
| 73 | return "" |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | func NormalizeQueueDrop(drop string) string { |
| 78 | switch strings.ToLower(strings.TrimSpace(drop)) { |
| 79 | case QueueDropOld: |
| 80 | return QueueDropOld |
| 81 | case QueueDropNew: |
| 82 | return QueueDropNew |
| 83 | default: |
| 84 | return QueueDropSummarize |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | // BuildSessionKey 根据 Hermes 模式生成稳定的 session key： |
| 89 | // - DM：按 chat 隔离（同一 DM 会话共享历史） |
| 90 | // - 群聊：按 user 隔离（每人独立会话） |
| 91 | // - thread：共享（thread 内所有人共享上下文） |
| 92 | func BuildSessionKey(src SessionSource) string { |
| 93 | var scope string |
| 94 | source := sessionSourceID(src) |
| 95 | switch src.ChatType { |
| 96 | case ChatDM: |
| 97 | scope = fmt.Sprintf("%s:dm:%s", source, src.ChatID) |
| 98 | case ChatGroup: |
| 99 | scope = fmt.Sprintf("%s:group:%s:%s", source, src.ChatID, src.UserID) |
| 100 | case ChatGuild: |
| 101 | scope = fmt.Sprintf("%s:guild:%s:%s", source, src.ChatID, src.UserID) |
| 102 | case ChatDirect: |
| 103 | scope = fmt.Sprintf("%s:direct:%s", source, src.ChatID) |
| 104 | case ChatThread: |
| 105 | threadID := src.ThreadID |
| 106 | if threadID == "" { |
| 107 | threadID = src.ChatID |
| 108 | } |
| 109 | scope = fmt.Sprintf("%s:thread:%s", source, threadID) |
| 110 | default: |
| 111 | scope = fmt.Sprintf("%s:%s:%s:%s", source, src.ChatType, src.ChatID, src.UserID) |
| 112 | } |
| 113 | h := sha256.Sum256([]byte(scope)) |
| 114 | return hex.EncodeToString(h[:])[:16] |
| 115 | } |
| 116 | |
| 117 | func sessionSourceID(src SessionSource) string { |
| 118 | if src.ConnectionID != "" { |
| 119 | return src.ConnectionID |
| 120 | } |
| 121 | if src.Domain != "" { |
| 122 | return fmt.Sprintf("%s:%s", src.Platform, src.Domain) |
| 123 | } |
| 124 | return string(src.Platform) |
| 125 | } |
| 126 | |
| 127 | // slashCommands 是绕过忙碌队列的命令集合。 |
| 128 | var slashCommands = map[string]bool{ |
| 129 | "/stop": true, |
| 130 | "/new": true, |
| 131 | "/reset": true, |
| 132 | "/approve": true, |
| 133 | "/deny": true, |
| 134 | "/answer": true, |
| 135 | "/yolo": true, |
| 136 | "/mode": true, |
| 137 | "/queue": true, |
| 138 | "/projects": true, |
| 139 | "/use": true, |
| 140 | "/sessions": true, |
| 141 | "/attach": true, |
| 142 | "/search": true, |
| 143 | "/desktop": true, |
| 144 | "/status": true, |
| 145 | "/help": true, |
| 146 | } |
| 147 | |
| 148 | // IsSlashBypass 判断消息是否为绕过队列的斜杠命令。 |
| 149 | func IsSlashBypass(text string) bool { |
| 150 | if len(text) == 0 { |
| 151 | return false |
| 152 | } |
| 153 | cmd := text |
| 154 | for i, r := range text { |
| 155 | if r == ' ' { |
| 156 | cmd = text[:i] |
| 157 | break |
| 158 | } |
| 159 | } |
| 160 | return slashCommands[cmd] |
| 161 | } |
| 162 | |
| 163 | // pendingTurn 是等待执行的一轮对话。 |
| 164 | type pendingTurn struct { |
| 165 | msg InboundMessage |
| 166 | timestamp time.Time |
| 167 | mode string |
| 168 | } |
| 169 | |
| 170 | // SessionManager 管理 session 级别的并发控制：同一 session 同时只跑一个任务。 |
| 171 | type SessionManager struct { |
| 172 | mu sync.Mutex |
| 173 | active map[string]bool // session key -> 是否正在运行 |
| 174 | pending map[string][]pendingTurn // session key -> 等待队列 |
| 175 | debounce time.Duration |
| 176 | modeOverrides map[string]string |
| 177 | dropped map[string][]string |
| 178 | } |
| 179 | |
| 180 | // NewSessionManager 创建一个新的 session 管理器。debounce 是消息合并窗口。 |
| 181 | func NewSessionManager(debounce time.Duration) *SessionManager { |
| 182 | if debounce <= 0 { |
| 183 | debounce = 1500 * time.Millisecond |
| 184 | } |
| 185 | return &SessionManager{ |
| 186 | active: make(map[string]bool), |
| 187 | pending: make(map[string][]pendingTurn), |
| 188 | debounce: debounce, |
| 189 | modeOverrides: make(map[string]string), |
| 190 | dropped: make(map[string][]string), |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | // TryAcquire 尝试获取 session 锁。如果 session 正忙且消息非绕过命令，返回 false。 |
| 195 | // 返回 (acquired, merged) — merged 为 true 表示消息已合并到等待队列。 |
| 196 | func (sm *SessionManager) TryAcquire(key string, msg InboundMessage) (acquired bool, merged bool) { |
| 197 | result := sm.TryAcquireWithQueue(key, msg, QueueOptions{Mode: QueueModeCollect, Cap: DefaultQueueCap, Drop: QueueDropSummarize}) |
| 198 | return result.Acquired, result.Queued |
| 199 | } |
| 200 | |
| 201 | func (sm *SessionManager) TryAcquireWithQueue(key string, msg InboundMessage, opts QueueOptions) QueueResult { |
| 202 | sm.mu.Lock() |
| 203 | defer sm.mu.Unlock() |
| 204 | |
| 205 | mode := NormalizeQueueMode(opts.Mode) |
| 206 | if mode == QueueModeSteer || mode == QueueModeInterrupt { |
| 207 | mode = QueueModeFollowup |
| 208 | } |
| 209 | cap := opts.Cap |
| 210 | if cap <= 0 { |
| 211 | cap = DefaultQueueCap |
| 212 | } |
| 213 | drop := NormalizeQueueDrop(opts.Drop) |
| 214 | |
| 215 | if sm.active[key] { |
| 216 | // 绕过命令立即返回 true（让调用方直接处理） |
| 217 | if IsSlashBypass(msg.Text) { |
| 218 | return QueueResult{Acquired: true, Mode: mode} |
| 219 | } |
| 220 | queue := sm.pending[key] |
| 221 | if len(queue) >= cap { |
| 222 | switch drop { |
| 223 | case QueueDropNew: |
| 224 | return QueueResult{Rejected: true, Pending: len(queue), Mode: mode} |
| 225 | case QueueDropOld, QueueDropSummarize: |
| 226 | removed := queue[0] |
| 227 | queue = queue[1:] |
| 228 | if drop == QueueDropSummarize { |
| 229 | sm.dropped[key] = append(sm.dropped[key], queueSummary(removed.msg.Text)) |
| 230 | } |
| 231 | } |
| 232 | } |
| 233 | if mode == QueueModeCollect && len(queue) > 0 { |
| 234 | last := &queue[len(queue)-1] |
| 235 | if msg.Text != "" && time.Since(last.timestamp) < sm.debounce { |
| 236 | if last.msg.Text != "" { |
| 237 | last.msg.Text = last.msg.Text + "\n" + msg.Text |
| 238 | } else { |
| 239 | last.msg.Text = msg.Text |
| 240 | } |
| 241 | last.timestamp = time.Now() |
| 242 | last.mode = mode |
| 243 | sm.pending[key] = queue |
| 244 | return QueueResult{Queued: true, Dropped: len(sm.dropped[key]) > 0, Pending: len(queue), Mode: mode} |
| 245 | } |
| 246 | } |
| 247 | queue = append(queue, pendingTurn{msg: msg, timestamp: time.Now(), mode: mode}) |
| 248 | sm.pending[key] = queue |
| 249 | return QueueResult{Queued: true, Dropped: len(sm.dropped[key]) > 0, Pending: len(queue), Mode: mode} |
| 250 | } |
| 251 | |
| 252 | sm.active[key] = true |
| 253 | return QueueResult{Acquired: true, Mode: mode} |
| 254 | } |
| 255 | |
| 256 | func (sm *SessionManager) ReplacePending(key string, msg InboundMessage) QueueResult { |
| 257 | sm.mu.Lock() |
| 258 | defer sm.mu.Unlock() |
| 259 | if !sm.active[key] { |
| 260 | sm.active[key] = true |
| 261 | return QueueResult{Acquired: true, Mode: QueueModeInterrupt} |
| 262 | } |
| 263 | sm.pending[key] = []pendingTurn{{msg: msg, timestamp: time.Now(), mode: QueueModeFollowup}} |
| 264 | delete(sm.dropped, key) |
| 265 | return QueueResult{Queued: true, Pending: 1, Mode: QueueModeInterrupt} |
| 266 | } |
| 267 | |
| 268 | // Release 释放 session 锁，返回等待队列中的下一条消息（合并后）。 |
| 269 | func (sm *SessionManager) Release(key string) *InboundMessage { |
| 270 | sm.mu.Lock() |
| 271 | defer sm.mu.Unlock() |
| 272 | |
| 273 | queue := sm.pending[key] |
| 274 | if len(queue) == 0 { |
| 275 | delete(sm.active, key) |
| 276 | delete(sm.pending, key) |
| 277 | delete(sm.dropped, key) |
| 278 | return nil |
| 279 | } |
| 280 | |
| 281 | mode := NormalizeQueueMode(queue[0].mode) |
| 282 | var merged *InboundMessage |
| 283 | if mode == QueueModeFollowup { |
| 284 | m := queue[0].msg |
| 285 | merged = &m |
| 286 | if len(queue) == 1 { |
| 287 | delete(sm.pending, key) |
| 288 | } else { |
| 289 | sm.pending[key] = queue[1:] |
| 290 | } |
| 291 | merged.Text = sm.consumeDroppedPrefixLocked(key, merged.Text) |
| 292 | return merged |
| 293 | } |
| 294 | |
| 295 | // collect 模式取出等待队列，并合并其中所有消息。 |
| 296 | for i := range queue { |
| 297 | if merged == nil { |
| 298 | m := queue[i].msg |
| 299 | merged = &m |
| 300 | } else { |
| 301 | if queue[i].msg.Text != "" { |
| 302 | merged.Text = merged.Text + "\n" + queue[i].msg.Text |
| 303 | } |
| 304 | } |
| 305 | } |
| 306 | delete(sm.pending, key) |
| 307 | merged.Text = sm.consumeDroppedPrefixLocked(key, merged.Text) |
| 308 | // active 保持 true，因为调用方会立即用 merged 消息开始新 turn |
| 309 | return merged |
| 310 | } |
| 311 | |
| 312 | func (sm *SessionManager) consumeDroppedPrefixLocked(key, text string) string { |
| 313 | dropped := sm.dropped[key] |
| 314 | if len(dropped) == 0 { |
| 315 | return text |
| 316 | } |
| 317 | delete(sm.dropped, key) |
| 318 | var b strings.Builder |
| 319 | fmt.Fprintf(&b, "[Queue note: %d older pending message(s) were dropped because this bot session reached its queue cap.", len(dropped)) |
| 320 | if len(dropped) > 0 { |
| 321 | b.WriteString(" Dropped summaries:") |
| 322 | limit := len(dropped) |
| 323 | if limit > 3 { |
| 324 | limit = 3 |
| 325 | } |
| 326 | for i := 0; i < limit; i++ { |
| 327 | fmt.Fprintf(&b, "\n- %s", dropped[i]) |
| 328 | } |
| 329 | if len(dropped) > limit { |
| 330 | fmt.Fprintf(&b, "\n- ... and %d more", len(dropped)-limit) |
| 331 | } |
| 332 | } |
| 333 | b.WriteString("]\n\n") |
| 334 | b.WriteString(text) |
| 335 | return b.String() |
| 336 | } |
| 337 | |
| 338 | func queueSummary(text string) string { |
| 339 | text = strings.TrimSpace(text) |
| 340 | if text == "" { |
| 341 | return "(empty message)" |
| 342 | } |
| 343 | runes := []rune(text) |
| 344 | if len(runes) <= 180 { |
| 345 | return text |
| 346 | } |
| 347 | return string(runes[:180]) + "..." |
| 348 | } |
| 349 | |
| 350 | // IsActive 返回 session 是否有正在运行的任务。 |
| 351 | func (sm *SessionManager) IsActive(key string) bool { |
| 352 | sm.mu.Lock() |
| 353 | defer sm.mu.Unlock() |
| 354 | return sm.active[key] |
| 355 | } |
| 356 | |
| 357 | // runIfIdle holds the per-gateway admission lock while fn switches runtime |
| 358 | // ownership for key. A normal message cannot become active between the idle |
| 359 | // check and the controller unlink/close sequence. |
| 360 | func (sm *SessionManager) runIfIdle(key string, fn func() bool) bool { |
| 361 | sm.mu.Lock() |
| 362 | defer sm.mu.Unlock() |
| 363 | if sm.active[key] || fn == nil { |
| 364 | return false |
| 365 | } |
| 366 | if !fn() { |
| 367 | return false |
| 368 | } |
| 369 | delete(sm.pending, key) |
| 370 | delete(sm.dropped, key) |
| 371 | return true |
| 372 | } |
| 373 | |
| 374 | // ActiveCount 返回当前活跃 session 数。 |
| 375 | func (sm *SessionManager) ActiveCount() int { |
| 376 | sm.mu.Lock() |
| 377 | defer sm.mu.Unlock() |
| 378 | return len(sm.active) |
| 379 | } |
| 380 | |
| 381 | func (sm *SessionManager) PendingCount(key string) int { |
| 382 | sm.mu.Lock() |
| 383 | defer sm.mu.Unlock() |
| 384 | return len(sm.pending[key]) |
| 385 | } |
| 386 | |
| 387 | func (sm *SessionManager) Snapshot() QueueSnapshot { |
| 388 | sm.mu.Lock() |
| 389 | defer sm.mu.Unlock() |
| 390 | var pending int |
| 391 | var dropped int |
| 392 | for _, queue := range sm.pending { |
| 393 | pending += len(queue) |
| 394 | } |
| 395 | for _, summaries := range sm.dropped { |
| 396 | dropped += len(summaries) |
| 397 | } |
| 398 | return QueueSnapshot{ |
| 399 | Active: len(sm.active), |
| 400 | Pending: pending, |
| 401 | Dropped: dropped, |
| 402 | Sessions: len(sm.active) + len(sm.pending), |
| 403 | } |
| 404 | } |
| 405 | |
| 406 | func (sm *SessionManager) QueueMode(key, fallback string) string { |
| 407 | sm.mu.Lock() |
| 408 | defer sm.mu.Unlock() |
| 409 | if mode := sm.modeOverrides[key]; mode != "" { |
| 410 | return mode |
| 411 | } |
| 412 | return NormalizeQueueMode(fallback) |
| 413 | } |
| 414 | |
| 415 | func (sm *SessionManager) SetQueueMode(key, mode string) { |
| 416 | sm.mu.Lock() |
| 417 | defer sm.mu.Unlock() |
| 418 | if normalized := NormalizeOptionalQueueMode(mode); normalized != "" { |
| 419 | sm.modeOverrides[key] = normalized |
| 420 | } |
| 421 | } |
| 422 | |
| 423 | func (sm *SessionManager) ClearQueueMode(key string) { |
| 424 | sm.mu.Lock() |
| 425 | defer sm.mu.Unlock() |
| 426 | delete(sm.modeOverrides, key) |
| 427 | } |
| 428 | |
| 429 | // ForceRelease 强制释放 session（用于 session 关闭或错误恢复）。 |
| 430 | func (sm *SessionManager) ForceRelease(key string) { |
| 431 | sm.mu.Lock() |
| 432 | defer sm.mu.Unlock() |
| 433 | delete(sm.active, key) |
| 434 | delete(sm.pending, key) |
| 435 | delete(sm.dropped, key) |
| 436 | } |
| 437 |