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:
viettranx committed 2026-02-24 17:44:50 +07:00
1 parent 16022d77be
commit f5ff96d998
10 files changed
+416 -85

No files matched your search

+5 -3
View File
@@ -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.
+8
View File
@@ -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
View File
@@ -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
View File
@@ -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))
+23 -9
View File
@@ -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
+1 -1
View File
@@ -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)
}
+4 -2
View File
@@ -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
+54
View File
@@ -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"},
+5 -3
View File
@@ -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
View File
@@ -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()