From e2debfe49a08894f42512086a4f0280e985e8876 Mon Sep 17 00:00:00 2001 From: viettranx Date: Tue, 3 Mar 2026 15:03:28 +0700 Subject: [PATCH] feat: mid-loop context compaction + team task user scoping Add mid-loop compaction to prevent context overflow during long-running delegated agent runs (e.g. 225K+ tokens causing DashScope timeouts). Uses same threshold as maybeSummarize (contextWindow * historyShare) with actual PromptTokens from LLM response. Only compacts the in-memory messages slice; pendingMsgs preserves full history for session flush. Add user_id/channel columns to team_tasks so end users only see their own tasks. Delegate/system channels bypass the filter to see all tasks. Group chats use the group-scoped UserID (group:channel:chatID) so all members share visibility. Co-Authored-By: Claude Opus 4.6 --- cmd/gateway_managed.go | 1 + internal/agent/loop.go | 122 +++++++++++++++++- internal/channels/telegram/commands_tasks.go | 6 +- internal/gateway/methods/teams.go | 2 +- internal/store/pg/teams_tasks.go | 36 ++++-- internal/store/team_store.go | 8 +- internal/tools/delegate.go | 3 +- internal/tools/team_tasks_tool.go | 20 ++- .../000008_team_tasks_user_scope.down.sql | 3 + .../000008_team_tasks_user_scope.up.sql | 3 + 10 files changed, 181 insertions(+), 23 deletions(-) create mode 100644 migrations/000008_team_tasks_user_scope.down.sql create mode 100644 migrations/000008_team_tasks_user_scope.up.sql diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index ad6905b9..4ddf7452 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -308,6 +308,7 @@ func wireManagedExtras( RunID: req.RunID, Stream: req.Stream, ExtraSystemPrompt: req.ExtraSystemPrompt, + MaxIterations: req.MaxIterations, }) if err != nil { return nil, err diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 9ce43d82..d57fcd45 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -260,6 +260,7 @@ type RunRequest struct { ParentRootSpanID uuid.UUID // if set, nest announce agent span under this parent span TraceName string // override trace name (default: "chat ") TraceTags []string // additional tags for the trace (e.g. "cron") + MaxIterations int // per-request override (0 = use agent default, must be lower) } // RunResult is the output of a completed agent run. @@ -547,6 +548,10 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) var mediaResults []MediaResult // media files from tool MEDIA: results var deliverables []string // actual content from tool outputs (for team task results) + // Mid-loop compaction: summarize in-memory messages when context exceeds threshold. + // Uses same config as maybeSummarize (contextWindow * historyShare). + var midLoopCompacted bool + // Team task orphan detection: track team_tasks create vs spawn calls. // If the LLM creates tasks but forgets to spawn, inject a reminder. var teamTaskCreates int // count of team_tasks action=create calls @@ -567,7 +572,12 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) }) }) - for iteration < l.maxIterations { + maxIter := l.maxIterations + if req.MaxIterations > 0 && req.MaxIterations < maxIter { + maxIter = req.MaxIterations + } + + for iteration < maxIter { iteration++ slog.Debug("agent iteration", "agent", l.id, "iteration", iteration, "messages", len(messages)) @@ -646,6 +656,36 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) totalUsage.ThinkingTokens += resp.Usage.ThinkingTokens } + // Mid-loop compaction: same threshold as maybeSummarize (contextWindow * historyShare) + // but applied to in-memory messages during the run. Prevents context overflow for + // long-running agents (e.g. delegated research tasks that accumulate many tool results). + if !midLoopCompacted && l.contextWindow > 0 { + historyShare := 0.75 + if l.compactionCfg != nil && l.compactionCfg.MaxHistoryShare > 0 { + historyShare = l.compactionCfg.MaxHistoryShare + } + threshold := int(float64(l.contextWindow) * historyShare) + + promptTokens := 0 + if resp.Usage != nil && resp.Usage.PromptTokens > 0 { + promptTokens = resp.Usage.PromptTokens + } else { + promptTokens = EstimateTokens(messages) + } + + if promptTokens >= threshold { + midLoopCompacted = true + if compacted := l.compactMessagesInPlace(ctx, messages); compacted != nil { + messages = compacted + } + slog.Info("mid_loop_compaction", + "agent", l.id, + "prompt_tokens", promptTokens, + "threshold", threshold, + "context_window", l.contextWindow) + } + } + // No tool calls → done if len(resp.ToolCalls) == 0 { // Guard: detect orphaned team_tasks create (created but not spawned). @@ -1000,6 +1040,86 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) }, nil } +// compactMessagesInPlace summarizes the first ~70% of messages into a condensed +// summary, keeping the last ~30% intact. Operates purely on the local messages +// slice — no session state touched, no locks needed. +// Returns nil on failure (caller keeps original messages). +func (l *Loop) compactMessagesInPlace(ctx context.Context, messages []providers.Message) []providers.Message { + if len(messages) < 6 { + return nil + } + + // Resolve keepCount from compaction config (same defaults as maybeSummarize). + keepCount := 4 + if l.compactionCfg != nil && l.compactionCfg.KeepLastMessages > 0 { + keepCount = l.compactionCfg.KeepLastMessages + } + // Ensure we keep at least 30% of messages. + if minKeep := len(messages) * 3 / 10; minKeep > keepCount { + keepCount = minKeep + } + + splitIdx := len(messages) - keepCount + + // Walk backward from splitIdx to find a clean boundary — + // avoid splitting tool_use → tool_result pairs. + for splitIdx > 0 { + m := messages[splitIdx] + if m.Role == "tool" || (m.Role == "assistant" && len(m.ToolCalls) > 0) { + splitIdx-- + continue + } + break + } + if splitIdx <= 1 { + return nil + } + + // Build summary input (same pattern as maybeSummarize in loop_history.go). + toSummarize := messages[:splitIdx] + var sb strings.Builder + for _, m := range toSummarize { + switch m.Role { + case "user": + fmt.Fprintf(&sb, "user: %s\n", m.Content) + case "assistant": + fmt.Fprintf(&sb, "assistant: %s\n", SanitizeAssistantContent(m.Content)) + } + } + + sctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + + resp, err := l.provider.Chat(sctx, providers.ChatRequest{ + Messages: []providers.Message{{ + Role: "user", + Content: "Provide a concise summary of this conversation, preserving key findings, data, and context:\n\n" + sb.String(), + }}, + Model: l.model, + Options: map[string]interface{}{"max_tokens": 1024, "temperature": 0.3}, + }) + if err != nil { + slog.Warn("mid_loop_compaction_failed", "agent", l.id, "error", err) + return nil + } + + summary := providers.Message{ + Role: "user", + Content: "[Summary of earlier conversation]\n" + SanitizeAssistantContent(resp.Content), + } + result := make([]providers.Message, 0, 1+keepCount) + result = append(result, summary) + result = append(result, messages[splitIdx:]...) + + slog.Info("mid_loop_compacted", + "agent", l.id, + "original_msgs", len(messages), + "summarized", splitIdx, + "kept", len(result)) + + return result +} + // parseMediaResult extracts a MediaResult from a tool result string containing "MEDIA:" prefix. // Handles formats: "MEDIA:/path/to/file" and "[[audio_as_voice]]\nMEDIA:/path/to/file". // Returns nil if no MEDIA: prefix is found. diff --git a/internal/channels/telegram/commands_tasks.go b/internal/channels/telegram/commands_tasks.go index 6f56dc32..f4254cbb 100644 --- a/internal/channels/telegram/commands_tasks.go +++ b/internal/channels/telegram/commands_tasks.go @@ -72,7 +72,7 @@ func (c *Channel) handleTasksList(ctx context.Context, chatID int64, setThread f return } - tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll) + tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll, "") if err != nil { slog.Warn("tasks command: ListTasks failed", "error", err) send("Failed to list tasks. Please try again.") @@ -164,7 +164,7 @@ func (c *Channel) handleTaskDetail(ctx context.Context, chatID int64, text strin return } - tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll) + tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll, "") if err != nil { slog.Warn("task_detail command: ListTasks failed", "error", err) send("Failed to list tasks. Please try again.") @@ -225,7 +225,7 @@ func (c *Channel) handleCallbackQuery(ctx context.Context, query *telego.Callbac return } - tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll) + tasks, err := c.teamStore.ListTasks(ctx, team.ID, "newest", store.TeamTaskFilterAll, "") if err != nil { send("Failed to list tasks.") return diff --git a/internal/gateway/methods/teams.go b/internal/gateway/methods/teams.go index 6451bbe5..69787be7 100644 --- a/internal/gateway/methods/teams.go +++ b/internal/gateway/methods/teams.go @@ -298,7 +298,7 @@ func (m *TeamsMethods) handleTaskList(_ context.Context, client *gateway.Client, } ctx := context.Background() - tasks, err := m.teamStore.ListTasks(ctx, teamID, "newest", store.TeamTaskFilterAll) + tasks, err := m.teamStore.ListTasks(ctx, teamID, "newest", store.TeamTaskFilterAll, "") if err != nil { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error())) return diff --git a/internal/store/pg/teams_tasks.go b/internal/store/pg/teams_tasks.go index e8631a07..7f65bbd9 100644 --- a/internal/store/pg/teams_tasks.go +++ b/internal/store/pg/teams_tasks.go @@ -25,11 +25,14 @@ func (s *PGTeamStore) CreateTask(ctx context.Context, task *store.TeamTaskData) task.UpdatedAt = now _, err := s.db.ExecContext(ctx, - `INSERT INTO team_tasks (id, team_id, subject, description, status, owner_agent_id, blocked_by, priority, result, created_at, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`, + `INSERT INTO team_tasks (id, team_id, subject, description, status, owner_agent_id, blocked_by, priority, result, user_id, channel, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`, task.ID, task.TeamID, task.Subject, task.Description, task.Status, task.OwnerAgentID, pq.Array(task.BlockedBy), - task.Priority, task.Result, now, now, + task.Priority, task.Result, + sql.NullString{String: task.UserID, Valid: task.UserID != ""}, + sql.NullString{String: task.Channel, Valid: task.Channel != ""}, + now, now, ) return err } @@ -42,7 +45,7 @@ func (s *PGTeamStore) UpdateTask(ctx context.Context, taskID uuid.UUID, updates return execMapUpdate(ctx, s.db, "team_tasks", taskID, updates) } -func (s *PGTeamStore) ListTasks(ctx context.Context, teamID uuid.UUID, orderBy string, statusFilter string) ([]store.TeamTaskData, error) { +func (s *PGTeamStore) ListTasks(ctx context.Context, teamID uuid.UUID, orderBy string, statusFilter string, userID string) ([]store.TeamTaskData, error) { orderClause := "t.priority DESC, t.created_at" if orderBy == "newest" { orderClause = "t.created_at DESC" @@ -57,12 +60,12 @@ func (s *PGTeamStore) ListTasks(ctx context.Context, teamID uuid.UUID, orderBy s } rows, err := s.db.QueryContext(ctx, - `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.created_at, t.updated_at, + `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.user_id, t.channel, t.created_at, t.updated_at, COALESCE(a.agent_key, '') AS owner_agent_key FROM team_tasks t LEFT JOIN agents a ON a.id = t.owner_agent_id - WHERE t.team_id = $1 `+statusWhere+` - ORDER BY `+orderClause, teamID) + WHERE t.team_id = $1 AND ($2 = '' OR t.user_id = $2) `+statusWhere+` + ORDER BY `+orderClause, teamID, userID) if err != nil { return nil, err } @@ -72,7 +75,7 @@ func (s *PGTeamStore) ListTasks(ctx context.Context, teamID uuid.UUID, orderBy s func (s *PGTeamStore) GetTask(ctx context.Context, taskID uuid.UUID) (*store.TeamTaskData, error) { rows, err := s.db.QueryContext(ctx, - `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.created_at, t.updated_at, + `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.user_id, t.channel, t.created_at, t.updated_at, COALESCE(a.agent_key, '') AS owner_agent_key FROM team_tasks t LEFT JOIN agents a ON a.id = t.owner_agent_id @@ -91,18 +94,18 @@ func (s *PGTeamStore) GetTask(ctx context.Context, taskID uuid.UUID) (*store.Tea return &tasks[0], nil } -func (s *PGTeamStore) SearchTasks(ctx context.Context, teamID uuid.UUID, query string, limit int) ([]store.TeamTaskData, error) { +func (s *PGTeamStore) SearchTasks(ctx context.Context, teamID uuid.UUID, query string, limit int, userID string) ([]store.TeamTaskData, error) { if limit <= 0 { limit = 20 } rows, err := s.db.QueryContext(ctx, - `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.created_at, t.updated_at, + `SELECT t.id, t.team_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.user_id, t.channel, t.created_at, t.updated_at, COALESCE(a.agent_key, '') AS owner_agent_key FROM team_tasks t LEFT JOIN agents a ON a.id = t.owner_agent_id - WHERE t.team_id = $1 AND t.tsv @@ plainto_tsquery('simple', $2) + WHERE t.team_id = $1 AND t.tsv @@ plainto_tsquery('simple', $2) AND ($4 = '' OR t.user_id = $4) ORDER BY ts_rank(t.tsv, plainto_tsquery('simple', $2)) DESC - LIMIT $3`, teamID, query, limit) + LIMIT $3`, teamID, query, limit, userID) if err != nil { return nil, err } @@ -173,12 +176,13 @@ func scanTaskRowsJoined(rows *sql.Rows) ([]store.TeamTaskData, error) { var tasks []store.TeamTaskData for rows.Next() { var d store.TeamTaskData - var desc, result sql.NullString + var desc, result, userID, channel sql.NullString var ownerID *uuid.UUID var blockedBy []uuid.UUID if err := rows.Scan( &d.ID, &d.TeamID, &d.Subject, &desc, &d.Status, &ownerID, pq.Array(&blockedBy), &d.Priority, &result, + &userID, &channel, &d.CreatedAt, &d.UpdatedAt, &d.OwnerAgentKey, ); err != nil { @@ -190,6 +194,12 @@ func scanTaskRowsJoined(rows *sql.Rows) ([]store.TeamTaskData, error) { if result.Valid { d.Result = &result.String } + if userID.Valid { + d.UserID = userID.String + } + if channel.Valid { + d.Channel = channel.String + } d.OwnerAgentID = ownerID d.BlockedBy = blockedBy tasks = append(tasks, d) diff --git a/internal/store/team_store.go b/internal/store/team_store.go index 14fb119d..060b6440 100644 --- a/internal/store/team_store.go +++ b/internal/store/team_store.go @@ -80,6 +80,8 @@ type TeamTaskData struct { Priority int `json:"priority"` Result *string `json:"result,omitempty"` Metadata map[string]interface{} `json:"metadata,omitempty"` + UserID string `json:"user_id,omitempty"` + Channel string `json:"channel,omitempty"` // Joined fields OwnerAgentKey string `json:"owner_agent_key,omitempty"` @@ -178,11 +180,13 @@ type TeamStore interface { UpdateTask(ctx context.Context, taskID uuid.UUID, updates map[string]any) error // ListTasks returns tasks for a team. orderBy: "priority" or "newest". // statusFilter: "" = non-completed (default), "completed", "all". - ListTasks(ctx context.Context, teamID uuid.UUID, orderBy string, statusFilter string) ([]TeamTaskData, error) + // userID: if non-empty, filter to tasks created by this user. + ListTasks(ctx context.Context, teamID uuid.UUID, orderBy string, statusFilter string, userID string) ([]TeamTaskData, error) // GetTask returns a single task by ID with joined agent info. GetTask(ctx context.Context, taskID uuid.UUID) (*TeamTaskData, error) // SearchTasks performs FTS search over task subject+description. - SearchTasks(ctx context.Context, teamID uuid.UUID, query string, limit int) ([]TeamTaskData, error) + // userID: if non-empty, filter to tasks created by this user. + SearchTasks(ctx context.Context, teamID uuid.UUID, query string, limit int, userID string) ([]TeamTaskData, error) // ClaimTask atomically transitions a task from pending to in_progress. // Only one agent can claim a given task (row-level lock, race-safe). diff --git a/internal/tools/delegate.go b/internal/tools/delegate.go index 103c4579..104518b1 100644 --- a/internal/tools/delegate.go +++ b/internal/tools/delegate.go @@ -71,6 +71,7 @@ type DelegateRunRequest struct { RunID string Stream bool ExtraSystemPrompt string + MaxIterations int // per-delegation override (0 = use agent default) } // DelegateRunResult is the result from AgentRunFunc. @@ -394,7 +395,7 @@ func (dm *DelegateManager) prepareDelegation(ctx context.Context, opts DelegateO // (common case: LLM called team_tasks create + spawn in parallel, // hallucinated the task_id, uuid.Parse failed → uuid.Nil). hint := "" - if tasks, err := dm.teamStore.ListTasks(ctx, team.ID, "newest", ""); err == nil { + if tasks, err := dm.teamStore.ListTasks(ctx, team.ID, "newest", "", ""); err == nil { var pendingIDs []string for _, t := range tasks { if t.Status == store.TeamTaskStatusPending { diff --git a/internal/tools/team_tasks_tool.go b/internal/tools/team_tasks_tool.go index 180e7458..af863808 100644 --- a/internal/tools/team_tasks_tool.go +++ b/internal/tools/team_tasks_tool.go @@ -104,7 +104,14 @@ func (t *TeamTasksTool) executeList(ctx context.Context, args map[string]interfa statusFilter, _ := args["status"].(string) - tasks, err := t.manager.teamStore.ListTasks(ctx, team.ID, "priority", statusFilter) + // Delegate/system channels see all tasks; end users only see their own. + filterUserID := "" + channel := ToolChannelFromCtx(ctx) + if channel != "delegate" && channel != "system" { + filterUserID = store.UserIDFromContext(ctx) + } + + tasks, err := t.manager.teamStore.ListTasks(ctx, team.ID, "priority", statusFilter, filterUserID) if err != nil { return ErrorResult("failed to list tasks: " + err.Error()) } @@ -180,7 +187,14 @@ func (t *TeamTasksTool) executeSearch(ctx context.Context, args map[string]inter return ErrorResult("query is required for search action") } - tasks, err := t.manager.teamStore.SearchTasks(ctx, team.ID, query, 20) + // Delegate/system channels see all tasks; end users only see their own. + filterUserID := "" + channel := ToolChannelFromCtx(ctx) + if channel != "delegate" && channel != "system" { + filterUserID = store.UserIDFromContext(ctx) + } + + tasks, err := t.manager.teamStore.SearchTasks(ctx, team.ID, query, 20, filterUserID) if err != nil { return ErrorResult("failed to search tasks: " + err.Error()) } @@ -244,6 +258,8 @@ func (t *TeamTasksTool) executeCreate(ctx context.Context, args map[string]inter Status: status, BlockedBy: blockedBy, Priority: priority, + UserID: store.UserIDFromContext(ctx), + Channel: ToolChannelFromCtx(ctx), } if err := t.manager.teamStore.CreateTask(ctx, task); err != nil { diff --git a/migrations/000008_team_tasks_user_scope.down.sql b/migrations/000008_team_tasks_user_scope.down.sql new file mode 100644 index 00000000..28371265 --- /dev/null +++ b/migrations/000008_team_tasks_user_scope.down.sql @@ -0,0 +1,3 @@ +DROP INDEX IF EXISTS idx_team_tasks_user_scope; +ALTER TABLE team_tasks DROP COLUMN IF EXISTS channel; +ALTER TABLE team_tasks DROP COLUMN IF EXISTS user_id; diff --git a/migrations/000008_team_tasks_user_scope.up.sql b/migrations/000008_team_tasks_user_scope.up.sql new file mode 100644 index 00000000..9607c6b7 --- /dev/null +++ b/migrations/000008_team_tasks_user_scope.up.sql @@ -0,0 +1,3 @@ +ALTER TABLE team_tasks ADD COLUMN user_id VARCHAR(255); +ALTER TABLE team_tasks ADD COLUMN channel VARCHAR(50); +CREATE INDEX idx_team_tasks_user_scope ON team_tasks(team_id, user_id) WHERE user_id IS NOT NULL;