diff --git a/README.md b/README.md index cee6d49e..8f8e5957 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/cmd/gateway.go b/cmd/gateway.go index 44df7e79..231cf6a6 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -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. diff --git a/cmd/gateway_consumer.go b/cmd/gateway_consumer.go index 555d4440..75c65da2 100644 --- a/cmd/gateway_consumer.go +++ b/cmd/gateway_consumer.go @@ -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) } diff --git a/internal/agent/loop.go b/internal/agent/loop.go index ba2b6019..8af1027e 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -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)) diff --git a/internal/agent/loop_history.go b/internal/agent/loop_history.go index 895a3da4..b8bc72c2 100644 --- a/internal/agent/loop_history.go +++ b/internal/agent/loop_history.go @@ -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 diff --git a/internal/agent/loop_history_test.go b/internal/agent/loop_history_test.go index a5ac98de..b44680f2 100644 --- a/internal/agent/loop_history_test.go +++ b/internal/agent/loop_history_test.go @@ -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) } diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index c774e95e..557a3782 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -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 diff --git a/internal/channels/telegram/commands.go b/internal/channels/telegram/commands.go index 4cc471f4..82423c71 100644 --- a/internal/channels/telegram/commands.go +++ b/internal/channels/telegram/commands.go @@ -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"}, diff --git a/internal/channels/telegram/reactions.go b/internal/channels/telegram/reactions.go index c063d539..fd7ad8f2 100644 --- a/internal/channels/telegram/reactions.go +++ b/internal/channels/telegram/reactions.go @@ -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() diff --git a/internal/scheduler/queue.go b/internal/scheduler/queue.go index 2941225e..194da628 100644 --- a/internal/scheduler/queue.go +++ b/internal/scheduler/queue.go @@ -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()