mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-17 04:15:05 +00:00
feat: Improve agent concurrency with per-session summarization locks and message buffering, add adaptive scheduler throttling, and introduce Telegram /stop and /stopall commands.
This commit is contained in:
1 parent
16022d77be
commit
f5ff96d998
10 files changed
+416
-85
No files matched your search
@@ -34,6 +34,8 @@ GoClaw is OpenClaw, reimagined in Go. It preserves the powerful gateway architec
|
||||
### Parallel Execution
|
||||
|
||||
- **Lane-based scheduler** — Main / subagent / cron lane isolation
|
||||
- **Concurrent agent runs** — Multiple users in a group chat get parallel responses (configurable `maxConcurrent`, default 3 for groups). Adaptive throttle reduces concurrency when session history nears summarization threshold
|
||||
- **Session history isolation** — Concurrent runs buffer messages locally and flush atomically on completion, preventing cross-run context pollution
|
||||
- **Concurrent subagents** — Depth limits, count limits, model override
|
||||
- **Batched announce queue** — Debounced result delivery
|
||||
|
||||
@@ -81,7 +83,7 @@ Standalone mode shares everything across users (same as OpenClaw). Managed mode
|
||||
| Security hardening | ✅ (SSRF, path traversal, injection) | ✅ (sandbox, rate limit, injection, pairing) | Basic (workspace restrict, exec deny) | ✅ 5-layer defense |
|
||||
| OTel observability | ✅ (opt-in extension) | ✅ (Prometheus + OTLP) | — | ✅ OTLP (opt-in build tag) |
|
||||
| Skill system | ✅ Embeddings/semantic | ✅ SKILL.md + TOML | ✅ Basic | ✅ BM25 + pgvector hybrid |
|
||||
| Lane-based scheduler | ✅ | Bounded concurrency | — | ✅ (main/subagent/cron) |
|
||||
| Lane-based scheduler | ✅ | Bounded concurrency | — | ✅ (main/subagent/cron + concurrent group runs) |
|
||||
| Messaging channels | 37+ | 15+ | 10+ | 5+ |
|
||||
| Companion apps | macOS, iOS, Android | Python SDK | — | Web dashboard |
|
||||
| Live Canvas / Voice | ✅ (A2UI + TTS/STT) | — | Voice transcription | TTS (4 providers) |
|
||||
@@ -96,7 +98,7 @@ Standalone mode shares everything across users (same as OpenClaw). Managed mode
|
||||
- **Multi-provider LLM support** — OpenRouter, Anthropic, OpenAI, Groq, DeepSeek, Gemini, Mistral, xAI, MiniMax, Cohere, Perplexity, and any OpenAI-compatible endpoint
|
||||
- **Agent loop** — Think-act-observe cycle with tool use, session history, and auto-summarization
|
||||
- **Subagents** — Spawn child agents with different models for parallel task execution
|
||||
- **Messaging channels** — Telegram, Discord, Zalo, Feishu/Lark, WhatsApp
|
||||
- **Messaging channels** — Telegram, Discord, Zalo, Feishu/Lark, WhatsApp. `/stop` cancels the current task, `/stopall` cancels all running tasks in a chat
|
||||
- **Memory system** — Long-term memory with SQLite FTS5 + vector embeddings (standalone) or pgvector hybrid search (managed)
|
||||
- **Skills** — SKILL.md-based knowledge base with BM25 search + embedding hybrid search (managed mode)
|
||||
- **Custom tools** — Define shell-based tools at runtime via HTTP API with JSON Schema parameters, auto shell-escaping, and encrypted environment variables (managed mode)
|
||||
@@ -636,7 +638,7 @@ GOCLAW_OPENROUTER_API_KEY=sk-or-xxx go test -v ./tests/integration/ -timeout 120
|
||||
- **WebSocket RPC protocol (v3)** — Connect handshake, chat streaming, event push all tested with web dashboard and integration tests.
|
||||
- **Store layer (PostgreSQL)** — All PG stores (sessions, agents, providers, skills, cron, pairing, tracing, memory) implemented and running in managed mode.
|
||||
- **Browser automation** — Rod/CDP integration for headless Chrome, tested in production agent workflows.
|
||||
- **Lane-based scheduler** — Main/subagent/cron lane isolation with concurrent execution tested.
|
||||
- **Lane-based scheduler** — Main/subagent/cron lane isolation with concurrent execution tested. Group chats support up to 3 concurrent agent runs per session with adaptive throttle and deferred session writes for history isolation.
|
||||
- **Security hardening** — Rate limiting, prompt injection detection, CORS, shell deny patterns, SSRF protection, credential scrubbing all implemented and verified.
|
||||
- **Web dashboard (core)** — Channel management, agent management, pairing approval, traces & spans viewer all implemented and working well.
|
||||
|
||||
|
||||
@@ -765,6 +765,14 @@ func runGateway() {
|
||||
)
|
||||
defer sched.Stop()
|
||||
|
||||
// Adaptive throttle: reduce per-session concurrency when nearing the summary threshold.
|
||||
// This prevents concurrent runs from racing with summarization.
|
||||
sched.SetTokenEstimateFunc(func(sessionKey string) (int, int) {
|
||||
history := sessStore.GetHistory(sessionKey)
|
||||
tokens := agent.EstimateTokens(history)
|
||||
return tokens, 200000 // default context window
|
||||
})
|
||||
|
||||
// Subscribe to agent events for channel streaming/reaction forwarding.
|
||||
// Events emitted by agent loops are broadcast to the bus; we forward them
|
||||
// to the channel manager which routes to StreamingChannel/ReactionChannel.
|
||||
|
||||
+59
-3
@@ -2,6 +2,7 @@ package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
@@ -97,9 +98,19 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
|
||||
)
|
||||
|
||||
// Enable streaming when the channel supports it (so agent emits chunk events).
|
||||
// Group chats: streaming disabled (concurrent runs would interleave chunks).
|
||||
enableStream := channelMgr != nil && channelMgr.IsStreamingChannel(msg.Channel)
|
||||
if peerKind == string(sessions.PeerGroup) {
|
||||
enableStream = false
|
||||
}
|
||||
|
||||
runID := fmt.Sprintf("inbound-%s-%s", msg.Channel, msg.ChatID)
|
||||
// Group chats allow concurrent runs (multiple users can chat simultaneously).
|
||||
maxConcurrent := 1
|
||||
if peerKind == string(sessions.PeerGroup) {
|
||||
maxConcurrent = 3
|
||||
}
|
||||
|
||||
runID := fmt.Sprintf("inbound-%s-%s-%s", msg.Channel, msg.ChatID, uuid.NewString()[:8])
|
||||
|
||||
// Register run with channel manager for streaming/reaction event forwarding.
|
||||
// Use localKey (composite key with topic suffix) so streaming/reaction events
|
||||
@@ -126,8 +137,8 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
|
||||
"- Address the group naturally. If the history shows a multi-person conversation, consider the full context before answering."
|
||||
}
|
||||
|
||||
// Schedule through main lane (per-session serialization + lane concurrency)
|
||||
outCh := sched.Schedule(ctx, "main", agent.RunRequest{
|
||||
// Schedule through main lane (per-session concurrency controlled by maxConcurrent)
|
||||
outCh := sched.ScheduleWithOpts(ctx, "main", agent.RunRequest{
|
||||
SessionKey: sessionKey,
|
||||
Message: msg.Content,
|
||||
Channel: msg.Channel,
|
||||
@@ -139,6 +150,8 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
|
||||
Stream: enableStream,
|
||||
HistoryLimit: msg.HistoryLimit,
|
||||
ExtraSystemPrompt: extraPrompt,
|
||||
}, scheduler.ScheduleOpts{
|
||||
MaxConcurrent: maxConcurrent,
|
||||
})
|
||||
|
||||
// Build outbound metadata for reply-to + thread routing.
|
||||
@@ -163,6 +176,18 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
|
||||
}
|
||||
|
||||
if outcome.Err != nil {
|
||||
// Don't send error for cancelled runs (/stop command) —
|
||||
// publish empty outbound to clean up thinking/typing indicators.
|
||||
if errors.Is(outcome.Err, context.Canceled) {
|
||||
slog.Info("inbound: run cancelled", "channel", channel, "session", session)
|
||||
msgBus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: channel,
|
||||
ChatID: chatID,
|
||||
Content: "",
|
||||
Metadata: meta,
|
||||
})
|
||||
return
|
||||
}
|
||||
slog.Error("inbound: agent run failed", "error", outcome.Err, "channel", channel)
|
||||
msgBus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: channel,
|
||||
@@ -317,6 +342,37 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
|
||||
continue
|
||||
}
|
||||
|
||||
// --- Command: /stop — cancel oldest active run for this session ---
|
||||
// --- Command: /stopall — cancel ALL active runs + drain queue ---
|
||||
if cmd := msg.Metadata["command"]; cmd == "stop" || cmd == "stopall" {
|
||||
agentID := msg.AgentID
|
||||
if agentID == "" {
|
||||
agentID = resolveAgentRoute(cfg, msg.Channel, msg.ChatID, msg.PeerKind)
|
||||
}
|
||||
peerKind := msg.PeerKind
|
||||
if peerKind == "" {
|
||||
peerKind = string(sessions.PeerDirect)
|
||||
}
|
||||
sessionKey := sessions.BuildScopedSessionKey(agentID, msg.Channel, sessions.PeerKind(peerKind), msg.ChatID, cfg.Sessions.Scope, cfg.Sessions.DmScope, cfg.Sessions.MainKey)
|
||||
if msg.Metadata["is_forum"] == "true" && peerKind == string(sessions.PeerGroup) {
|
||||
var topicID int
|
||||
fmt.Sscanf(msg.Metadata["message_thread_id"], "%d", &topicID)
|
||||
if topicID > 0 {
|
||||
sessionKey = sessions.BuildGroupTopicSessionKey(agentID, msg.Channel, msg.ChatID, topicID)
|
||||
}
|
||||
}
|
||||
|
||||
var cancelled bool
|
||||
if cmd == "stopall" {
|
||||
cancelled = sched.CancelSession(sessionKey)
|
||||
slog.Info("inbound: /stopall command", "session", sessionKey, "cancelled", cancelled)
|
||||
} else {
|
||||
cancelled = sched.CancelOneSession(sessionKey)
|
||||
slog.Info("inbound: /stop command", "session", sessionKey, "cancelled", cancelled)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// --- Normal messages: route through debouncer ---
|
||||
debouncer.Push(msg)
|
||||
}
|
||||
|
||||
+29
-11
@@ -48,8 +48,10 @@ type Loop struct {
|
||||
tools *tools.Registry
|
||||
toolPolicy *tools.PolicyEngine // optional: filters tools sent to LLM
|
||||
agentToolPolicy *config.ToolPolicySpec // per-agent tool policy from DB (nil = no restrictions)
|
||||
running atomic.Bool
|
||||
mu sync.Mutex // protects concurrent runs
|
||||
activeRuns atomic.Int32 // number of currently executing runs
|
||||
|
||||
// Per-session summarization lock: prevents concurrent summarize goroutines for the same session.
|
||||
summarizeMu sync.Map // sessionKey → *sync.Mutex
|
||||
|
||||
// Bootstrap/persona context (loaded at startup, injected into system prompt)
|
||||
ownerIDs []string
|
||||
@@ -229,8 +231,8 @@ type RunResult struct {
|
||||
// Run processes a single message through the agent loop.
|
||||
// It blocks until completion and returns the final response.
|
||||
func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
l.activeRuns.Add(1)
|
||||
defer l.activeRuns.Add(-1)
|
||||
|
||||
l.emit(AgentEvent{Type: "run.started", AgentID: l.id, RunID: req.RunID})
|
||||
|
||||
@@ -294,8 +296,16 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) {
|
||||
Payload: map[string]string{"error": err.Error()},
|
||||
})
|
||||
// Only finish trace for root runs; child traces don't own the trace lifecycle.
|
||||
// Use background context when the run context is cancelled (/stop command)
|
||||
// so the DB update still succeeds.
|
||||
if !isChildTrace && l.traceCollector != nil && traceID != uuid.Nil {
|
||||
l.traceCollector.FinishTrace(ctx, traceID, "error", err.Error(), "")
|
||||
traceCtx := ctx
|
||||
traceStatus := "error"
|
||||
if ctx.Err() != nil {
|
||||
traceCtx = context.Background()
|
||||
traceStatus = "cancelled"
|
||||
}
|
||||
l.traceCollector.FinishTrace(traceCtx, traceID, traceStatus, err.Error(), "")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
@@ -399,8 +409,10 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error)
|
||||
|
||||
messages := l.buildMessages(ctx, history, summary, req.Message, req.ExtraSystemPrompt, req.SessionKey, req.Channel, req.UserID, req.HistoryLimit)
|
||||
|
||||
// 2. Save user message to session
|
||||
l.sessions.AddMessage(req.SessionKey, providers.Message{
|
||||
// 2. Buffer new messages — write to session only AFTER the run completes.
|
||||
// This prevents concurrent runs from seeing each other's in-progress messages.
|
||||
var pendingMsgs []providers.Message
|
||||
pendingMsgs = append(pendingMsgs, providers.Message{
|
||||
Role: "user",
|
||||
Content: req.Message,
|
||||
})
|
||||
@@ -481,7 +493,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error)
|
||||
ToolCalls: resp.ToolCalls,
|
||||
}
|
||||
messages = append(messages, assistantMsg)
|
||||
l.sessions.AddMessage(req.SessionKey, assistantMsg)
|
||||
pendingMsgs = append(pendingMsgs, assistantMsg)
|
||||
|
||||
// Execute tool calls (parallel when multiple, sequential when single)
|
||||
if len(resp.ToolCalls) == 1 {
|
||||
@@ -531,7 +543,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error)
|
||||
ToolCallID: tc.ID,
|
||||
}
|
||||
messages = append(messages, toolMsg)
|
||||
l.sessions.AddMessage(req.SessionKey, toolMsg)
|
||||
pendingMsgs = append(pendingMsgs, toolMsg)
|
||||
} else {
|
||||
// Multiple tools: parallel execution via goroutines.
|
||||
// Tool instances are immutable (context-based) so concurrent access is safe.
|
||||
@@ -617,7 +629,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error)
|
||||
ToolCallID: r.tc.ID,
|
||||
}
|
||||
messages = append(messages, toolMsg)
|
||||
l.sessions.AddMessage(req.SessionKey, toolMsg)
|
||||
pendingMsgs = append(pendingMsgs, toolMsg)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -639,11 +651,17 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error)
|
||||
}
|
||||
}
|
||||
|
||||
l.sessions.AddMessage(req.SessionKey, providers.Message{
|
||||
pendingMsgs = append(pendingMsgs, providers.Message{
|
||||
Role: "assistant",
|
||||
Content: finalContent,
|
||||
})
|
||||
|
||||
// Flush all buffered messages to session atomically.
|
||||
// This ensures concurrent runs never see each other's in-progress messages.
|
||||
for _, msg := range pendingMsgs {
|
||||
l.sessions.AddMessage(req.SessionKey, msg)
|
||||
}
|
||||
|
||||
// Write session metadata (matching TS session entry updates)
|
||||
l.sessions.UpdateMetadata(req.SessionKey, l.model, l.provider.Name(), req.Channel)
|
||||
l.sessions.AccumulateTokens(req.SessionKey, int64(totalUsage.PromptTokens), int64(totalUsage.CompletionTokens))
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -230,32 +231,45 @@ func sanitizeHistory(msgs []providers.Message) []providers.Message {
|
||||
|
||||
func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) {
|
||||
history := l.sessions.GetHistory(sessionKey)
|
||||
tokenEstimate := estimateTokens(history)
|
||||
tokenEstimate := EstimateTokens(history)
|
||||
threshold := l.contextWindow * 75 / 100
|
||||
|
||||
if len(history) <= 50 && tokenEstimate <= threshold {
|
||||
return
|
||||
}
|
||||
|
||||
// Memory flush: run BEFORE compaction so agent can save important context.
|
||||
// Matching TS: memory flush is a separate embedded agent turn before summarization.
|
||||
// Per-session lock: prevent concurrent summarize+flush goroutines for the same session.
|
||||
// TryLock is non-blocking — if another run is already summarizing this session, skip.
|
||||
// The next run will trigger summarization again if still needed.
|
||||
muI, _ := l.summarizeMu.LoadOrStore(sessionKey, &sync.Mutex{})
|
||||
sessionMu := muI.(*sync.Mutex)
|
||||
if !sessionMu.TryLock() {
|
||||
slog.Debug("summarization already in progress, skipping", "session", sessionKey)
|
||||
return
|
||||
}
|
||||
|
||||
// Memory flush runs synchronously INSIDE the guard
|
||||
// (so concurrent runs don't both trigger flush for the same compaction cycle).
|
||||
flushSettings := ResolveMemoryFlushSettings(l.compactionCfg)
|
||||
if l.shouldRunMemoryFlush(sessionKey, tokenEstimate, flushSettings) {
|
||||
// Run flush synchronously before compaction (so writes happen before truncation)
|
||||
l.runMemoryFlush(ctx, sessionKey, flushSettings)
|
||||
}
|
||||
|
||||
// Summarize in background
|
||||
// Summarize in background (holds the per-session lock until done)
|
||||
go func() {
|
||||
sctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
|
||||
defer cancel()
|
||||
|
||||
summary := l.sessions.GetSummary(sessionKey)
|
||||
defer sessionMu.Unlock()
|
||||
|
||||
// Re-check: history may have been truncated by a concurrent summarize
|
||||
// that finished between our threshold check and acquiring the lock.
|
||||
history := l.sessions.GetHistory(sessionKey)
|
||||
if len(history) <= 4 {
|
||||
return
|
||||
}
|
||||
|
||||
sctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
|
||||
defer cancel()
|
||||
|
||||
summary := l.sessions.GetSummary(sessionKey)
|
||||
toSummarize := history[:len(history)-4]
|
||||
|
||||
var sb string
|
||||
|
||||
@@ -209,7 +209,7 @@ func TestEstimateTokens(t *testing.T) {
|
||||
{Role: "user", Content: "Hello world!"}, // 12 chars → ~4 tokens
|
||||
{Role: "assistant", Content: "Hi there, how are you?"}, // 22 chars → ~7 tokens
|
||||
}
|
||||
got := estimateTokens(msgs)
|
||||
got := EstimateTokens(msgs)
|
||||
if got <= 0 {
|
||||
t.Errorf("expected positive token estimate, got %d", got)
|
||||
}
|
||||
|
||||
@@ -28,7 +28,7 @@ func (l *Loop) ID() string { return l.id }
|
||||
func (l *Loop) Model() string { return l.model }
|
||||
|
||||
// IsRunning returns whether the agent is currently processing.
|
||||
func (l *Loop) IsRunning() bool { return l.running.Load() }
|
||||
func (l *Loop) IsRunning() bool { return l.activeRuns.Load() > 0 }
|
||||
|
||||
// emitLLMSpan records an LLM call span if tracing is active.
|
||||
// When GOCLAW_TRACE_VERBOSE is set, messages are serialized as InputPreview.
|
||||
@@ -184,7 +184,9 @@ func truncateStr(s string, maxLen int) string {
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
|
||||
func estimateTokens(messages []providers.Message) int {
|
||||
// EstimateTokens returns a rough token estimate for a slice of messages.
|
||||
// Used internally for summarization thresholds and externally for adaptive throttle.
|
||||
func EstimateTokens(messages []providers.Message) int {
|
||||
total := 0
|
||||
for _, m := range messages {
|
||||
total += utf8.RuneCountInString(m.Content) / 3
|
||||
|
||||
@@ -67,6 +67,8 @@ func (c *Channel) handleBotCommand(ctx context.Context, message *telego.Message,
|
||||
helpText := "Available commands:\n" +
|
||||
"/start — Start chatting with the bot\n" +
|
||||
"/help — Show this help message\n" +
|
||||
"/stop — Stop current running task\n" +
|
||||
"/stopall — Stop all running tasks\n" +
|
||||
"/reset — Reset conversation history\n" +
|
||||
"/status — Show bot status\n" +
|
||||
"/writers — List file writers for this group\n" +
|
||||
@@ -104,6 +106,56 @@ func (c *Channel) handleBotCommand(ctx context.Context, message *telego.Message,
|
||||
c.bot.SendMessage(ctx, msg)
|
||||
return true
|
||||
|
||||
case "/stop":
|
||||
peerKind := "direct"
|
||||
if isGroup {
|
||||
peerKind = "group"
|
||||
}
|
||||
c.Bus().PublishInbound(bus.InboundMessage{
|
||||
Channel: c.Name(),
|
||||
SenderID: senderID,
|
||||
ChatID: chatIDStr,
|
||||
Content: "/stop",
|
||||
PeerKind: peerKind,
|
||||
AgentID: c.AgentID(),
|
||||
UserID: strings.SplitN(senderID, "|", 2)[0],
|
||||
Metadata: map[string]string{
|
||||
"command": "stop",
|
||||
"local_key": localKey,
|
||||
"is_forum": fmt.Sprintf("%t", isForum),
|
||||
"message_thread_id": fmt.Sprintf("%d", messageThreadID),
|
||||
},
|
||||
})
|
||||
msg := tu.Message(chatIDObj, "Stopping current task…")
|
||||
setThread(msg)
|
||||
c.bot.SendMessage(ctx, msg)
|
||||
return true
|
||||
|
||||
case "/stopall":
|
||||
peerKind := "direct"
|
||||
if isGroup {
|
||||
peerKind = "group"
|
||||
}
|
||||
c.Bus().PublishInbound(bus.InboundMessage{
|
||||
Channel: c.Name(),
|
||||
SenderID: senderID,
|
||||
ChatID: chatIDStr,
|
||||
Content: "/stopall",
|
||||
PeerKind: peerKind,
|
||||
AgentID: c.AgentID(),
|
||||
UserID: strings.SplitN(senderID, "|", 2)[0],
|
||||
Metadata: map[string]string{
|
||||
"command": "stopall",
|
||||
"local_key": localKey,
|
||||
"is_forum": fmt.Sprintf("%t", isForum),
|
||||
"message_thread_id": fmt.Sprintf("%d", messageThreadID),
|
||||
},
|
||||
})
|
||||
msg := tu.Message(chatIDObj, "Stopping all tasks…")
|
||||
setThread(msg)
|
||||
c.bot.SendMessage(ctx, msg)
|
||||
return true
|
||||
|
||||
case "/status":
|
||||
statusText := fmt.Sprintf("Bot status: Running\nChannel: Telegram\nBot: @%s", c.bot.Username())
|
||||
msg := tu.Message(chatIDObj, statusText)
|
||||
@@ -376,6 +428,8 @@ func DefaultMenuCommands() []telego.BotCommand {
|
||||
return []telego.BotCommand{
|
||||
{Command: "start", Description: "Start chatting with the bot"},
|
||||
{Command: "help", Description: "Show available commands"},
|
||||
{Command: "stop", Description: "Stop current running task"},
|
||||
{Command: "stopall", Description: "Stop all running tasks"},
|
||||
{Command: "reset", Description: "Reset conversation history"},
|
||||
{Command: "status", Description: "Show bot status"},
|
||||
{Command: "writers", Description: "List file writers for this group"},
|
||||
|
||||
@@ -2,6 +2,7 @@ package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -247,8 +248,9 @@ func (c *Channel) OnReactionEvent(ctx context.Context, chatID string, messageID
|
||||
return err
|
||||
}
|
||||
|
||||
// Get or create reaction controller for this message
|
||||
key := chatID
|
||||
// Get or create reaction controller for this message.
|
||||
// Key by messageID so concurrent runs in the same chat don't clash.
|
||||
key := fmt.Sprintf("%s:%d", chatID, messageID)
|
||||
val, _ := c.reactions.LoadOrStore(key, newStatusReactionController(c.bot, id, messageID))
|
||||
rc := val.(*StatusReactionController)
|
||||
|
||||
@@ -270,7 +272,7 @@ func (c *Channel) ClearReaction(ctx context.Context, chatID string, messageID in
|
||||
}
|
||||
|
||||
// Stop and remove controller if exists
|
||||
key := chatID
|
||||
key := fmt.Sprintf("%s:%d", chatID, messageID)
|
||||
if val, ok := c.reactions.LoadAndDelete(key); ok {
|
||||
rc := val.(*StatusReactionController)
|
||||
rc.Stop()
|
||||
|
||||
+228
-53
@@ -34,19 +34,21 @@ const (
|
||||
|
||||
// QueueConfig configures per-session message queuing.
|
||||
type QueueConfig struct {
|
||||
Mode QueueMode `json:"mode"`
|
||||
Cap int `json:"cap"`
|
||||
Drop DropPolicy `json:"drop"`
|
||||
DebounceMs int `json:"debounce_ms"`
|
||||
Mode QueueMode `json:"mode"`
|
||||
Cap int `json:"cap"`
|
||||
Drop DropPolicy `json:"drop"`
|
||||
DebounceMs int `json:"debounce_ms"`
|
||||
MaxConcurrent int `json:"max_concurrent"` // 0 or 1 = serial (default)
|
||||
}
|
||||
|
||||
// DefaultQueueConfig returns sensible defaults.
|
||||
func DefaultQueueConfig() QueueConfig {
|
||||
return QueueConfig{
|
||||
Mode: QueueModeQueue,
|
||||
Cap: 10,
|
||||
Drop: DropOld,
|
||||
DebounceMs: 800,
|
||||
Mode: QueueModeQueue,
|
||||
Cap: 10,
|
||||
Drop: DropOld,
|
||||
DebounceMs: 800,
|
||||
MaxConcurrent: 1,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,6 +56,10 @@ func DefaultQueueConfig() QueueConfig {
|
||||
// The scheduler calls this when it's the request's turn.
|
||||
type RunFunc func(ctx context.Context, req agent.RunRequest) (*agent.RunResult, error)
|
||||
|
||||
// TokenEstimateFunc returns token estimate and context window for a session.
|
||||
// Used by adaptive throttle to reduce concurrency near the summary threshold.
|
||||
type TokenEstimateFunc func(sessionKey string) (tokens int, contextWindow int)
|
||||
|
||||
// PendingRequest is a queued agent run awaiting execution.
|
||||
type PendingRequest struct {
|
||||
Req agent.RunRequest
|
||||
@@ -66,8 +72,8 @@ type RunOutcome struct {
|
||||
Err error
|
||||
}
|
||||
|
||||
// SessionQueue serializes agent runs for a single session key.
|
||||
// Only one run executes at a time; additional messages are queued.
|
||||
// SessionQueue manages agent runs for a single session key.
|
||||
// Supports configurable concurrency: 1 (serial) or N (concurrent).
|
||||
type SessionQueue struct {
|
||||
key string
|
||||
config QueueConfig
|
||||
@@ -75,27 +81,71 @@ type SessionQueue struct {
|
||||
laneMgr *LaneManager
|
||||
lane string
|
||||
|
||||
mu sync.Mutex
|
||||
queue []*PendingRequest
|
||||
active bool // whether a run is currently executing
|
||||
cancel context.CancelFunc // cancel for the active run (interrupt mode)
|
||||
timer *time.Timer // debounce timer
|
||||
parentCtx context.Context // stored from first Enqueue call, used for spawning runs
|
||||
mu sync.Mutex
|
||||
queue []*PendingRequest
|
||||
activeRuns map[string]context.CancelFunc // runID → cancel
|
||||
activeOrder []string // FIFO order of active runIDs
|
||||
maxConcurrent int // effective limit (from config or per-session override)
|
||||
timer *time.Timer // debounce timer
|
||||
parentCtx context.Context // stored from first Enqueue call
|
||||
|
||||
tokenEstimateFn TokenEstimateFunc // optional: for adaptive throttle
|
||||
}
|
||||
|
||||
// NewSessionQueue creates a queue for a specific session.
|
||||
func NewSessionQueue(key, lane string, cfg QueueConfig, laneMgr *LaneManager, runFn RunFunc) *SessionQueue {
|
||||
maxC := cfg.MaxConcurrent
|
||||
if maxC <= 0 {
|
||||
maxC = 1
|
||||
}
|
||||
return &SessionQueue{
|
||||
key: key,
|
||||
config: cfg,
|
||||
runFn: runFn,
|
||||
laneMgr: laneMgr,
|
||||
lane: lane,
|
||||
key: key,
|
||||
config: cfg,
|
||||
runFn: runFn,
|
||||
laneMgr: laneMgr,
|
||||
lane: lane,
|
||||
activeRuns: make(map[string]context.CancelFunc),
|
||||
maxConcurrent: maxC,
|
||||
}
|
||||
}
|
||||
|
||||
// SetMaxConcurrent overrides the per-session max concurrent runs.
|
||||
// Typically called from the consumer when it knows the peer kind (group vs DM).
|
||||
func (sq *SessionQueue) SetMaxConcurrent(n int) {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
sq.maxConcurrent = n
|
||||
}
|
||||
|
||||
// effectiveMaxConcurrent returns the current concurrency limit,
|
||||
// reduced to 1 when near the summary threshold (adaptive throttle).
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) effectiveMaxConcurrent() int {
|
||||
max := sq.maxConcurrent
|
||||
if max <= 0 {
|
||||
max = 1
|
||||
}
|
||||
if sq.tokenEstimateFn == nil {
|
||||
return max
|
||||
}
|
||||
tokens, contextWindow := sq.tokenEstimateFn(sq.key)
|
||||
if contextWindow > 0 && float64(tokens)/float64(contextWindow) >= 0.6 {
|
||||
return 1 // near summary threshold → serialize
|
||||
}
|
||||
return max
|
||||
}
|
||||
|
||||
// hasCapacity returns whether a new run can start.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) hasCapacity() bool {
|
||||
return len(sq.activeRuns) < sq.effectiveMaxConcurrent()
|
||||
}
|
||||
|
||||
// Enqueue adds a request to the session queue.
|
||||
// If no run is active, it starts immediately (after debounce).
|
||||
// If capacity is available, it starts immediately (after debounce).
|
||||
// Returns a channel that receives the result when the run completes.
|
||||
func (sq *SessionQueue) Enqueue(ctx context.Context, req agent.RunRequest) <-chan RunOutcome {
|
||||
outcome := make(chan RunOutcome, 1)
|
||||
@@ -111,14 +161,16 @@ func (sq *SessionQueue) Enqueue(ctx context.Context, req agent.RunRequest) <-cha
|
||||
|
||||
switch sq.config.Mode {
|
||||
case QueueModeInterrupt:
|
||||
// Cancel current run if active
|
||||
if sq.active && sq.cancel != nil {
|
||||
sq.cancel()
|
||||
// Cancel all active runs
|
||||
for runID, cancel := range sq.activeRuns {
|
||||
cancel()
|
||||
delete(sq.activeRuns, runID)
|
||||
}
|
||||
sq.activeOrder = nil
|
||||
// Clear existing queue and enqueue this one
|
||||
sq.drainQueue(RunOutcome{Err: context.Canceled})
|
||||
sq.queue = append(sq.queue, pending)
|
||||
if !sq.active {
|
||||
if sq.hasCapacity() {
|
||||
sq.scheduleNext(ctx)
|
||||
}
|
||||
|
||||
@@ -129,7 +181,7 @@ func (sq *SessionQueue) Enqueue(ctx context.Context, req agent.RunRequest) <-cha
|
||||
sq.queue = append(sq.queue, pending)
|
||||
}
|
||||
|
||||
if !sq.active {
|
||||
if sq.hasCapacity() {
|
||||
sq.scheduleNext(ctx)
|
||||
}
|
||||
}
|
||||
@@ -137,7 +189,7 @@ func (sq *SessionQueue) Enqueue(ctx context.Context, req agent.RunRequest) <-cha
|
||||
return outcome
|
||||
}
|
||||
|
||||
// scheduleNext starts the next queued request, applying debounce.
|
||||
// scheduleNext starts the next queued request(s), applying debounce.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) scheduleNext(ctx context.Context) {
|
||||
if len(sq.queue) == 0 {
|
||||
@@ -146,7 +198,7 @@ func (sq *SessionQueue) scheduleNext(ctx context.Context) {
|
||||
|
||||
debounce := time.Duration(sq.config.DebounceMs) * time.Millisecond
|
||||
if debounce <= 0 {
|
||||
sq.startNext(ctx)
|
||||
sq.startAvailable(ctx)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -157,25 +209,34 @@ func (sq *SessionQueue) scheduleNext(ctx context.Context) {
|
||||
sq.timer = time.AfterFunc(debounce, func() {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
if !sq.active && len(sq.queue) > 0 {
|
||||
sq.startNext(ctx)
|
||||
if sq.hasCapacity() && len(sq.queue) > 0 {
|
||||
sq.startAvailable(ctx)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// startNext picks the first queued request and runs it in the lane.
|
||||
// startAvailable starts as many queued requests as capacity allows.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) startNext(ctx context.Context) {
|
||||
func (sq *SessionQueue) startAvailable(ctx context.Context) {
|
||||
for sq.hasCapacity() && len(sq.queue) > 0 {
|
||||
sq.startOne(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// startOne picks the first queued request and runs it in the lane.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) startOne(ctx context.Context) {
|
||||
if len(sq.queue) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
pending := sq.queue[0]
|
||||
sq.queue = sq.queue[1:]
|
||||
sq.active = true
|
||||
|
||||
runID := pending.Req.RunID
|
||||
runCtx, cancel := context.WithCancel(ctx)
|
||||
sq.cancel = cancel
|
||||
sq.activeRuns[runID] = cancel
|
||||
sq.activeOrder = append(sq.activeOrder, runID)
|
||||
|
||||
lane := sq.laneMgr.Get(sq.lane)
|
||||
if lane == nil {
|
||||
@@ -184,39 +245,50 @@ func (sq *SessionQueue) startNext(ctx context.Context) {
|
||||
|
||||
if lane == nil {
|
||||
// No lane available — run directly
|
||||
go sq.executeRun(runCtx, pending)
|
||||
go sq.executeRun(runCtx, runID, pending)
|
||||
return
|
||||
}
|
||||
|
||||
err := lane.Submit(ctx, func() {
|
||||
sq.executeRun(runCtx, pending)
|
||||
sq.executeRun(runCtx, runID, pending)
|
||||
})
|
||||
if err != nil {
|
||||
pending.ResultCh <- RunOutcome{Err: err}
|
||||
close(pending.ResultCh)
|
||||
// caller already holds sq.mu — set fields directly
|
||||
sq.active = false
|
||||
sq.cancel = nil
|
||||
// caller already holds sq.mu — clean up
|
||||
delete(sq.activeRuns, runID)
|
||||
sq.removeFromOrder(runID)
|
||||
}
|
||||
}
|
||||
|
||||
// executeRun runs the agent and then processes the next queued message.
|
||||
func (sq *SessionQueue) executeRun(ctx context.Context, pending *PendingRequest) {
|
||||
// executeRun runs the agent and then starts the next queued message(s) if capacity allows.
|
||||
func (sq *SessionQueue) executeRun(ctx context.Context, runID string, pending *PendingRequest) {
|
||||
result, err := sq.runFn(ctx, pending.Req)
|
||||
pending.ResultCh <- RunOutcome{Result: result, Err: err}
|
||||
close(pending.ResultCh)
|
||||
|
||||
sq.mu.Lock()
|
||||
sq.active = false
|
||||
sq.cancel = nil
|
||||
delete(sq.activeRuns, runID)
|
||||
sq.removeFromOrder(runID)
|
||||
|
||||
if len(sq.queue) > 0 {
|
||||
// Use parentCtx (not the per-run ctx which may be cancelled in interrupt mode)
|
||||
if sq.hasCapacity() && len(sq.queue) > 0 {
|
||||
// Use parentCtx (not the per-run ctx which may be cancelled)
|
||||
sq.scheduleNext(sq.parentCtx)
|
||||
}
|
||||
sq.mu.Unlock()
|
||||
}
|
||||
|
||||
// removeFromOrder removes a runID from the activeOrder slice.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) removeFromOrder(runID string) {
|
||||
for i, id := range sq.activeOrder {
|
||||
if id == runID {
|
||||
sq.activeOrder = append(sq.activeOrder[:i], sq.activeOrder[i+1:]...)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyDropPolicy handles a full queue.
|
||||
// Must be called with sq.mu held.
|
||||
func (sq *SessionQueue) applyDropPolicy(incoming *PendingRequest) {
|
||||
@@ -258,11 +330,63 @@ func (sq *SessionQueue) drainQueue(outcome RunOutcome) {
|
||||
sq.queue = nil
|
||||
}
|
||||
|
||||
// IsActive returns whether a run is currently executing.
|
||||
// CancelOne stops the oldest active run (FIFO).
|
||||
// Does NOT drain the pending queue. Used by /stop command.
|
||||
// Returns true if an active run was actually cancelled.
|
||||
func (sq *SessionQueue) CancelOne() bool {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
|
||||
if len(sq.activeOrder) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
// Cancel the oldest active run
|
||||
runID := sq.activeOrder[0]
|
||||
if cancel, ok := sq.activeRuns[runID]; ok {
|
||||
cancel()
|
||||
delete(sq.activeRuns, runID)
|
||||
sq.activeOrder = sq.activeOrder[1:]
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// CancelAll stops all active runs and drains all pending requests.
|
||||
// Used by /stopall command.
|
||||
// Returns true if any active run was actually cancelled.
|
||||
func (sq *SessionQueue) CancelAll() bool {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
|
||||
cancelled := false
|
||||
for runID, cancel := range sq.activeRuns {
|
||||
cancel()
|
||||
delete(sq.activeRuns, runID)
|
||||
cancelled = true
|
||||
}
|
||||
sq.activeOrder = nil
|
||||
sq.drainQueue(RunOutcome{Err: context.Canceled})
|
||||
return cancelled
|
||||
}
|
||||
|
||||
// Cancel is an alias for CancelAll (backward compat with /stop command).
|
||||
func (sq *SessionQueue) Cancel() bool {
|
||||
return sq.CancelAll()
|
||||
}
|
||||
|
||||
// IsActive returns whether any run is currently executing.
|
||||
func (sq *SessionQueue) IsActive() bool {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
return sq.active
|
||||
return len(sq.activeRuns) > 0
|
||||
}
|
||||
|
||||
// ActiveCount returns the number of currently executing runs.
|
||||
func (sq *SessionQueue) ActiveCount() int {
|
||||
sq.mu.Lock()
|
||||
defer sq.mu.Unlock()
|
||||
return len(sq.activeRuns)
|
||||
}
|
||||
|
||||
// QueueLen returns the number of pending messages.
|
||||
@@ -272,13 +396,21 @@ func (sq *SessionQueue) QueueLen() int {
|
||||
return len(sq.queue)
|
||||
}
|
||||
|
||||
// --- Scheduler ---
|
||||
|
||||
// ScheduleOpts provides per-request overrides for the scheduler.
|
||||
type ScheduleOpts struct {
|
||||
MaxConcurrent int // per-session override (0 = use config default)
|
||||
}
|
||||
|
||||
// Scheduler is the top-level coordinator that manages lanes and session queues.
|
||||
type Scheduler struct {
|
||||
lanes *LaneManager
|
||||
sessions map[string]*SessionQueue
|
||||
config QueueConfig
|
||||
runFn RunFunc
|
||||
mu sync.RWMutex
|
||||
lanes *LaneManager
|
||||
sessions map[string]*SessionQueue
|
||||
config QueueConfig
|
||||
runFn RunFunc
|
||||
mu sync.RWMutex
|
||||
tokenEstimateFn TokenEstimateFunc // optional: for adaptive throttle
|
||||
}
|
||||
|
||||
// NewScheduler creates a scheduler with the given lane and queue config.
|
||||
@@ -295,6 +427,12 @@ func NewScheduler(laneConfigs []LaneConfig, queueCfg QueueConfig, runFn RunFunc)
|
||||
}
|
||||
}
|
||||
|
||||
// SetTokenEstimateFunc sets the callback used by adaptive throttle.
|
||||
// Must be called before any Schedule calls.
|
||||
func (s *Scheduler) SetTokenEstimateFunc(fn TokenEstimateFunc) {
|
||||
s.tokenEstimateFn = fn
|
||||
}
|
||||
|
||||
// Schedule submits a run request to the appropriate session queue and lane.
|
||||
// Returns a channel that receives the result when the run completes.
|
||||
func (s *Scheduler) Schedule(ctx context.Context, lane string, req agent.RunRequest) <-chan RunOutcome {
|
||||
@@ -302,6 +440,15 @@ func (s *Scheduler) Schedule(ctx context.Context, lane string, req agent.RunRequ
|
||||
return sq.Enqueue(ctx, req)
|
||||
}
|
||||
|
||||
// ScheduleWithOpts submits a run request with per-session overrides.
|
||||
func (s *Scheduler) ScheduleWithOpts(ctx context.Context, lane string, req agent.RunRequest, opts ScheduleOpts) <-chan RunOutcome {
|
||||
sq := s.getOrCreateSession(req.SessionKey, lane)
|
||||
if opts.MaxConcurrent > 0 {
|
||||
sq.SetMaxConcurrent(opts.MaxConcurrent)
|
||||
}
|
||||
return sq.Enqueue(ctx, req)
|
||||
}
|
||||
|
||||
// getOrCreateSession returns or creates a session queue for the given key.
|
||||
func (s *Scheduler) getOrCreateSession(sessionKey, lane string) *SessionQueue {
|
||||
s.mu.RLock()
|
||||
@@ -321,12 +468,40 @@ func (s *Scheduler) getOrCreateSession(sessionKey, lane string) *SessionQueue {
|
||||
}
|
||||
|
||||
sq = NewSessionQueue(sessionKey, lane, s.config, s.lanes, s.runFn)
|
||||
if s.tokenEstimateFn != nil {
|
||||
sq.tokenEstimateFn = s.tokenEstimateFn
|
||||
}
|
||||
s.sessions[sessionKey] = sq
|
||||
|
||||
slog.Debug("session queue created", "session", sessionKey, "lane", lane)
|
||||
return sq
|
||||
}
|
||||
|
||||
// CancelSession cancels all active runs and drains pending queue for a session.
|
||||
// Returns true if any active run was cancelled.
|
||||
func (s *Scheduler) CancelSession(sessionKey string) bool {
|
||||
s.mu.RLock()
|
||||
sq, ok := s.sessions[sessionKey]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return sq.CancelAll()
|
||||
}
|
||||
|
||||
// CancelOneSession cancels the oldest active run for a session.
|
||||
// Does NOT drain the pending queue. Used by /stop command.
|
||||
// Returns true if an active run was cancelled.
|
||||
func (s *Scheduler) CancelOneSession(sessionKey string) bool {
|
||||
s.mu.RLock()
|
||||
sq, ok := s.sessions[sessionKey]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return sq.CancelOne()
|
||||
}
|
||||
|
||||
// Stop shuts down all lanes and clears session queues.
|
||||
func (s *Scheduler) Stop() {
|
||||
s.lanes.StopAll()
|
||||
|
||||
Reference in new issue
Block a user