From 619b253b822527c0c5d04767a27f2fc528d3b981 Mon Sep 17 00:00:00 2001 From: viettranx Date: Tue, 14 Apr 2026 10:05:57 +0700 Subject: [PATCH 1/3] fix(tasks): inject tenant ctx in task ticker to prevent nil panic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: ticker's recoverCtx had no tenant → PGTeamStore.GetTeam returned silent (nil, nil) → team.LeadAgentID nil-deref panic - Fix notifyLeaders: composite cache keys {TeamID, TenantID}, inject scopeCtx = store.WithTenantID(ctx, scope.TenantID) before GetTeam/GetByID/GetTask, nil-check team + lead agent - Fix processFollowups: per-team scopeCtx from teamTasks[0].TenantID, nil-check team before followupInterval(*team) - Add TenantID field to TeamTaskData + scan paths in PG and SQLite stores - Bonus: GetTask(scopeCtx, ...) propagates tenant for peerKind session routing (related #266) - Tests: upgrade stub to function-based dispatch + ctx capture, add 6 regression tests (nil-team no-panic, multi-tenant cache isolation, cache hit dedup, multi-tenant ctx in processFollowups) - Docs: scheduling-cron guide notes tenant-ctx injection requirement for background workers --- docs/08-scheduling-cron.md | 1 + internal/store/pg/teams_tasks.go | 4 +- internal/store/sqlitestore/teams_tasks.go | 4 +- internal/store/team_store.go | 1 + internal/tasks/task_ticker.go | 69 +++++- internal/tasks/task_ticker_test.go | 274 +++++++++++++++++++++- 6 files changed, 331 insertions(+), 22 deletions(-) diff --git a/docs/08-scheduling-cron.md b/docs/08-scheduling-cron.md index 7f23b5b0..d62de27a 100644 --- a/docs/08-scheduling-cron.md +++ b/docs/08-scheduling-cron.md @@ -122,6 +122,7 @@ Cancel commands for Telegram and other channels. - **Empty outbound**: On cancel, an empty outbound message is published to trigger cleanup (stop typing indicator, clear reactions) - **Trace finalization**: When `ctx.Err() != nil`, trace finalization falls back to `context.Background()` for the final DB write. Status is set to `"cancelled"` - **Context survival**: Context values (traceID, collector) survive cancellation -- only the Done channel fires +- **Background workers (ticker/cron) — tenant ctx injection required**: Jobs started from `context.Background()` carry no tenant. Before calling any tenant-scoped store method (e.g. `GetTeam`, `GetTask`, `GetByID`), the worker MUST inject `store.WithTenantID(ctx, tenantID)` derived from the row-level `tenant_id` (e.g. `RecoveredTaskInfo.TenantID`, `TeamTaskData.TenantID`). Callers must also nil-check returned entities — some stores (e.g. `PGTeamStore.GetTeam`) return `(nil, nil)` when tenant is missing rather than an error. See `internal/tasks/task_ticker.go` for the reference pattern - **Generation counter**: Each `SessionQueue` tracks a generation counter. When reset (e.g., during SIGUSR1 in-process restart), old generations are ignored, preventing stale completions from interfering with new requests --- diff --git a/internal/store/pg/teams_tasks.go b/internal/store/pg/teams_tasks.go index dac6e2c9..4b825df1 100644 --- a/internal/store/pg/teams_tasks.go +++ b/internal/store/pg/teams_tasks.go @@ -22,7 +22,7 @@ import ( const taskLockDuration = 60 * time.Minute // taskSelectCols is the shared SELECT column list for task queries (must match scanTaskRowsJoined). -const taskSelectCols = `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, +const taskSelectCols = `t.id, t.team_id, t.tenant_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.user_id, t.channel, t.task_type, t.task_number, COALESCE(t.identifier,''), t.created_by_agent_id, COALESCE(t.assignee_user_id,''), t.parent_id, COALESCE(t.chat_id,''), t.metadata, t.locked_at, t.lock_expires_at, COALESCE(t.progress_percent,0), COALESCE(t.progress_step,''), t.followup_at, COALESCE(t.followup_count,0), COALESCE(t.followup_max,0), COALESCE(t.followup_message,''), COALESCE(t.followup_channel,''), COALESCE(t.followup_chat_id,''), @@ -560,7 +560,7 @@ func scanTaskRowsJoined(rows *sql.Rows) ([]store.TeamTaskData, error) { var followupCount, followupMax int var followupMessage, followupChannel, followupChatID string if err := rows.Scan( - &d.ID, &d.TeamID, &d.Subject, &desc, &d.Status, + &d.ID, &d.TeamID, &d.TenantID, &d.Subject, &desc, &d.Status, &ownerID, pq.Array(&blockedBy), &d.Priority, &result, &userID, &channel, &d.TaskType, &d.TaskNumber, &identifier, &createdByAgentID, &assigneeUserID, &parentID, diff --git a/internal/store/sqlitestore/teams_tasks.go b/internal/store/sqlitestore/teams_tasks.go index 0699666f..85540e4b 100644 --- a/internal/store/sqlitestore/teams_tasks.go +++ b/internal/store/sqlitestore/teams_tasks.go @@ -24,7 +24,7 @@ const taskLockDuration = 30 * time.Minute const maxListTasksRows = 30 // taskSelectCols is the shared SELECT column list for task queries. -const taskSelectCols = `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, +const taskSelectCols = `t.id, t.team_id, t.tenant_id, t.subject, t.description, t.status, t.owner_agent_id, t.blocked_by, t.priority, t.result, t.user_id, t.channel, t.task_type, t.task_number, COALESCE(t.identifier,''), t.created_by_agent_id, COALESCE(t.assignee_user_id,''), t.parent_id, COALESCE(t.chat_id,''), t.metadata, t.locked_at, t.lock_expires_at, COALESCE(t.progress_percent,0), COALESCE(t.progress_step,''), t.followup_at, COALESCE(t.followup_count,0), COALESCE(t.followup_max,0), COALESCE(t.followup_message,''), COALESCE(t.followup_channel,''), COALESCE(t.followup_chat_id,''), @@ -510,7 +510,7 @@ func scanTaskRowsJoined(rows *sql.Rows) ([]store.TeamTaskData, error) { var followupMessage, followupChannel, followupChatID string createdAt, updatedAt := scanTimePair() if err := rows.Scan( - &d.ID, &d.TeamID, &d.Subject, &desc, &d.Status, + &d.ID, &d.TeamID, &d.TenantID, &d.Subject, &desc, &d.Status, &ownerID, &blockedByJSON, &d.Priority, &result, &userID, &channel, &d.TaskType, &d.TaskNumber, &identifier, &createdByAgentID, &assigneeUserID, &parentID, diff --git a/internal/store/team_store.go b/internal/store/team_store.go index 61ec76f7..d954530e 100644 --- a/internal/store/team_store.go +++ b/internal/store/team_store.go @@ -94,6 +94,7 @@ type TeamMemberData struct { type TeamTaskData struct { BaseModel TeamID uuid.UUID `json:"team_id" db:"team_id"` + TenantID uuid.UUID `json:"tenant_id" db:"tenant_id"` Subject string `json:"subject" db:"subject"` Description string `json:"description,omitempty" db:"description"` Status string `json:"status" db:"status"` diff --git a/internal/tasks/task_ticker.go b/internal/tasks/task_ticker.go index fbe85509..e5775317 100644 --- a/internal/tasks/task_ticker.go +++ b/internal/tasks/task_ticker.go @@ -195,27 +195,59 @@ func (t *TaskTicker) notifyLeaders(ctx context.Context, tasks []store.RecoveredT } // Cache team+lead lookups (same team may have multiple scopes). - teamCache := map[uuid.UUID]*store.TeamData{} - leadCache := map[uuid.UUID]*store.AgentData{} + // Composite key includes TenantID to clarify multi-tenant intent (UUIDs are globally + // unique but composite key makes the isolation boundary explicit and future-proof). + type teamCacheKey struct { + TeamID uuid.UUID + TenantID uuid.UUID + } + type leadCacheKey struct { + AgentID uuid.UUID + TenantID uuid.UUID + } + teamCache := map[teamCacheKey]*store.TeamData{} + leadCache := map[leadCacheKey]*store.AgentData{} for scope, scopeTasks := range byScope { - team := teamCache[scope.TeamID] + // Inject tenant into ctx so store lookups (GetTeam, GetByID, GetTask) apply the + // correct tenant filter. Without this, PGTeamStore.GetTeam silently returns + // (nil, nil) when ctx has no tenant — causing a nil-deref panic on team.LeadAgentID. + scopeCtx := ctx + if scope.TenantID != uuid.Nil { + scopeCtx = store.WithTenantID(ctx, scope.TenantID) + } + + teamKey := teamCacheKey{TeamID: scope.TeamID, TenantID: scope.TenantID} + team := teamCache[teamKey] if team == nil { var err error - team, err = t.teams.GetTeam(ctx, scope.TeamID) + team, err = t.teams.GetTeam(scopeCtx, scope.TeamID) if err != nil { + slog.Warn("task_ticker: get team failed", "team_id", scope.TeamID, "tenant_id", scope.TenantID, "error", err) continue } - teamCache[scope.TeamID] = team + if team == nil { + // GetTeam can return (nil, nil) when tenant ctx is missing or team is deleted. + slog.Warn("task_ticker: team not found (nil)", "team_id", scope.TeamID, "tenant_id", scope.TenantID) + continue + } + teamCache[teamKey] = team } - lead := leadCache[team.LeadAgentID] + + leadKey := leadCacheKey{AgentID: team.LeadAgentID, TenantID: scope.TenantID} + lead := leadCache[leadKey] if lead == nil { var err error - lead, err = t.agents.GetByID(ctx, team.LeadAgentID) + lead, err = t.agents.GetByID(scopeCtx, team.LeadAgentID) if err != nil { + slog.Warn("task_ticker: get lead agent failed", "agent_id", team.LeadAgentID, "tenant_id", scope.TenantID, "error", err) continue } - leadCache[team.LeadAgentID] = lead + if lead == nil { + slog.Warn("task_ticker: lead agent not found (nil)", "agent_id", team.LeadAgentID, "tenant_id", scope.TenantID) + continue + } + leadCache[leadKey] = lead } // Build batched task list with clear actionable hints. @@ -238,7 +270,7 @@ func (t *TaskTicker) notifyLeaders(ctx context.Context, tasks []store.RecoveredT // Resolve PeerKind from first task's metadata for correct session routing (#266). var peerKind string var fullTask *store.TeamTaskData - if task, err := t.teams.GetTask(ctx, scopeTasks[0].ID); err == nil { + if task, err := t.teams.GetTask(scopeCtx, scopeTasks[0].ID); err == nil { fullTask = task if fullTask != nil && fullTask.Metadata != nil { if pk, ok := fullTask.Metadata["peer_kind"].(string); ok { @@ -304,12 +336,27 @@ func (t *TaskTicker) processFollowups(ctx context.Context) { byTeam[task.TeamID] = append(byTeam[task.TeamID], task) } for teamID, teamTasks := range byTeam { - team, err := t.teams.GetTeam(ctx, teamID) + if len(teamTasks) == 0 { + continue + } + tenantID := teamTasks[0].TenantID + scopeCtx := ctx + if tenantID != uuid.Nil { + scopeCtx = store.WithTenantID(ctx, tenantID) + } + team, err := t.teams.GetTeam(scopeCtx, teamID) if err != nil { + slog.Warn("task_ticker: followups get team failed", + "team_id", teamID, "tenant_id", tenantID, "error", err) + continue + } + if team == nil { + slog.Warn("task_ticker: followups team not found (nil)", + "team_id", teamID, "tenant_id", tenantID) continue } interval := followupInterval(*team) - t.processTeamFollowups(ctx, teamTasks, interval) + t.processTeamFollowups(scopeCtx, teamTasks, interval) } } diff --git a/internal/tasks/task_ticker_test.go b/internal/tasks/task_ticker_test.go index d93802f1..432c6ca5 100644 --- a/internal/tasks/task_ticker_test.go +++ b/internal/tasks/task_ticker_test.go @@ -18,6 +18,12 @@ import ( // ─── minimal stub stores ─────────────────────────────────────────────────── +// getTeamCall records ctx metadata for each GetTeam call (tenant propagation assertions). +type getTeamCall struct { + TeamID uuid.UUID + TenantID uuid.UUID +} + // stubTeamStore satisfies store.TeamStore by embedding the interface and only // overriding the methods actually exercised by TaskTicker. type stubTeamStore struct { @@ -37,11 +43,18 @@ type stubTeamStore struct { recoverErr error staleErr error - recovered []store.RecoveredTaskInfo - stale []store.RecoveredTaskInfo - inReview []store.RecoveredTaskInfo - orphans []store.RecoveredTaskInfo + recovered []store.RecoveredTaskInfo + stale []store.RecoveredTaskInfo + inReview []store.RecoveredTaskInfo + orphans []store.RecoveredTaskInfo followupTasks []store.TeamTaskData + + // function-based dispatch — nil means default behavior (return error "not found") + getTeamFunc func(ctx context.Context, id uuid.UUID) (*store.TeamData, error) + getTaskFunc func(ctx context.Context, id uuid.UUID) (*store.TeamTaskData, error) + + // captured calls for tenant propagation assertions (protected by mu) + getTeamCalls []getTeamCall } func (s *stubTeamStore) ForceRecoverAllTasks(_ context.Context) ([]store.RecoveredTaskInfo, error) { @@ -79,17 +92,38 @@ func (s *stubTeamStore) IncrementFollowupCount(_ context.Context, _ uuid.UUID, _ return nil } -func (s *stubTeamStore) GetTeam(_ context.Context, _ uuid.UUID) (*store.TeamData, error) { +func (s *stubTeamStore) GetTeam(ctx context.Context, id uuid.UUID) (*store.TeamData, error) { + s.mu.Lock() + s.getTeamCalls = append(s.getTeamCalls, getTeamCall{ + TeamID: id, + TenantID: store.TenantIDFromContext(ctx), + }) + s.mu.Unlock() + if s.getTeamFunc != nil { + return s.getTeamFunc(ctx, id) + } return nil, errors.New("not found") } -func (s *stubTeamStore) GetTask(_ context.Context, _ uuid.UUID) (*store.TeamTaskData, error) { +func (s *stubTeamStore) GetTask(ctx context.Context, id uuid.UUID) (*store.TeamTaskData, error) { + if s.getTaskFunc != nil { + return s.getTaskFunc(ctx, id) + } return nil, errors.New("not found") } -// stubAgentStore satisfies store.AgentStore — unused by ticker logic. +// stubAgentStore satisfies store.AgentStore — unused by most ticker logic. +// GetByID is overridable via getByIDFunc for tests exercising the lead-agent lookup path. type stubAgentStore struct { store.AgentStore + getByIDFunc func(ctx context.Context, id uuid.UUID) (*store.AgentData, error) +} + +func (s *stubAgentStore) GetByID(ctx context.Context, id uuid.UUID) (*store.AgentData, error) { + if s.getByIDFunc != nil { + return s.getByIDFunc(ctx, id) + } + return nil, errors.New("not found") } // ─── NewTaskTicker ───────────────────────────────────────────────────────── @@ -491,3 +525,229 @@ func TestFollowupOutboundMessage_OmitsLocalKeyWhenMissing(t *testing.T) { t.Fatalf("expected metadata to be nil, got %#v", got.Metadata) } } + +// ─── notifyLeaders: nil-team and tenant propagation ────────────────────────── + +// TestNotifyLeaders_NilTeam_NoPanic verifies that when GetTeam returns (nil, nil) +// the scope is skipped gracefully — no panic, no publishes. +func TestNotifyLeaders_NilTeam_NoPanic(t *testing.T) { + tenantID := uuid.New() + teamID := uuid.New() + + ts := &stubTeamStore{ + getTeamFunc: func(_ context.Context, _ uuid.UUID) (*store.TeamData, error) { + return nil, nil // simulate PGTeamStore returning nil without error + }, + } + as := &stubAgentStore{} + mb := bus.New() + tt := NewTaskTicker(ts, as, mb, 3600) + + // Must not panic. + tt.notifyLeaders(context.Background(), []store.RecoveredTaskInfo{ + {ID: uuid.New(), TeamID: teamID, TenantID: tenantID}, + }, "recovered", "hint") + + ts.mu.Lock() + callCount := len(ts.getTeamCalls) + capturedTenant := uuid.Nil + if callCount > 0 { + capturedTenant = ts.getTeamCalls[0].TenantID + } + ts.mu.Unlock() + + if callCount != 1 { + t.Errorf("GetTeam call count = %d, want 1", callCount) + } + if capturedTenant != tenantID { + t.Errorf("GetTeam ctx TenantID = %v, want %v", capturedTenant, tenantID) + } +} + +// TestNotifyLeaders_MultiTenantCacheIsolation verifies that tasks from different tenants +// each trigger their own GetTeam call with the correct tenant in ctx. +func TestNotifyLeaders_MultiTenantCacheIsolation(t *testing.T) { + tenantX := uuid.New() + tenantY := uuid.New() + teamT1 := uuid.New() + teamT2 := uuid.New() + leadID := uuid.New() + + ts := &stubTeamStore{ + getTeamFunc: func(_ context.Context, id uuid.UUID) (*store.TeamData, error) { + return &store.TeamData{ + BaseModel: store.BaseModel{ID: id}, + LeadAgentID: leadID, + }, nil + }, + } + as := &stubAgentStore{ + // short-circuit before publish; return error so notifyLeaders skips publishing + getByIDFunc: func(_ context.Context, _ uuid.UUID) (*store.AgentData, error) { + return nil, errors.New("not found") + }, + } + mb := bus.New() + tt := NewTaskTicker(ts, as, mb, 3600) + + tasks := []store.RecoveredTaskInfo{ + {ID: uuid.New(), TeamID: teamT1, TenantID: tenantX}, + {ID: uuid.New(), TeamID: teamT2, TenantID: tenantY}, + } + tt.notifyLeaders(context.Background(), tasks, "recovered", "hint") + + ts.mu.Lock() + calls := make([]getTeamCall, len(ts.getTeamCalls)) + copy(calls, ts.getTeamCalls) + ts.mu.Unlock() + + if len(calls) != 2 { + t.Fatalf("GetTeam call count = %d, want 2", len(calls)) + } + + // Build map: teamID → tenantID from captured calls. + captured := map[uuid.UUID]uuid.UUID{} + for _, c := range calls { + captured[c.TeamID] = c.TenantID + } + + if captured[teamT1] != tenantX { + t.Errorf("team T1: ctx TenantID = %v, want %v", captured[teamT1], tenantX) + } + if captured[teamT2] != tenantY { + t.Errorf("team T2: ctx TenantID = %v, want %v", captured[teamT2], tenantY) + } +} + +// TestNotifyLeaders_CachedOnSameTenant verifies that two scopes with the same team+tenant +// result in only one GetTeam call (cache hit on composite key). +func TestNotifyLeaders_CachedOnSameTenant(t *testing.T) { + tenantID := uuid.New() + teamID := uuid.New() + leadID := uuid.New() + + ts := &stubTeamStore{ + getTeamFunc: func(_ context.Context, id uuid.UUID) (*store.TeamData, error) { + return &store.TeamData{ + BaseModel: store.BaseModel{ID: id}, + LeadAgentID: leadID, + }, nil + }, + } + as := &stubAgentStore{ + getByIDFunc: func(_ context.Context, _ uuid.UUID) (*store.AgentData, error) { + return nil, errors.New("not found") + }, + } + mb := bus.New() + tt := NewTaskTicker(ts, as, mb, 3600) + + // Two tasks same team+tenant but different ChatID → two scopes, same cache key. + tasks := []store.RecoveredTaskInfo{ + {ID: uuid.New(), TeamID: teamID, TenantID: tenantID, ChatID: "chat-A"}, + {ID: uuid.New(), TeamID: teamID, TenantID: tenantID, ChatID: "chat-B"}, + } + tt.notifyLeaders(context.Background(), tasks, "recovered", "hint") + + ts.mu.Lock() + callCount := len(ts.getTeamCalls) + ts.mu.Unlock() + + if callCount != 1 { + t.Errorf("GetTeam call count = %d, want 1 (cache hit)", callCount) + } +} + +// ─── processFollowups: multi-tenant tenant ctx and nil-team guard ───────────── + +// TestProcessFollowups_MultiTenantTenantCtx verifies that tasks from two different +// tenants each trigger GetTeam with the correct tenant in ctx. +func TestProcessFollowups_MultiTenantTenantCtx(t *testing.T) { + tenantX := uuid.New() + tenantY := uuid.New() + teamT1 := uuid.New() + teamT2 := uuid.New() + + ts := &stubTeamStore{ + followupTasks: []store.TeamTaskData{ + { + BaseModel: store.BaseModel{ID: uuid.New()}, + TeamID: teamT1, + TenantID: tenantX, + FollowupChannel: "telegram", + FollowupChatID: "chat-1", + FollowupMessage: "ping", + }, + { + BaseModel: store.BaseModel{ID: uuid.New()}, + TeamID: teamT2, + TenantID: tenantY, + FollowupChannel: "telegram", + FollowupChatID: "chat-2", + FollowupMessage: "ping", + }, + }, + getTeamFunc: func(_ context.Context, id uuid.UUID) (*store.TeamData, error) { + return &store.TeamData{BaseModel: store.BaseModel{ID: id}}, nil + }, + } + as := &stubAgentStore{} + mb := bus.New() + tt := NewTaskTicker(ts, as, mb, 3600) + + tt.processFollowups(context.Background()) + + ts.mu.Lock() + calls := make([]getTeamCall, len(ts.getTeamCalls)) + copy(calls, ts.getTeamCalls) + ts.mu.Unlock() + + if len(calls) != 2 { + t.Fatalf("GetTeam call count = %d, want 2", len(calls)) + } + + captured := map[uuid.UUID]uuid.UUID{} + for _, c := range calls { + captured[c.TeamID] = c.TenantID + } + + if captured[teamT1] != tenantX { + t.Errorf("team T1: ctx TenantID = %v, want %v", captured[teamT1], tenantX) + } + if captured[teamT2] != tenantY { + t.Errorf("team T2: ctx TenantID = %v, want %v", captured[teamT2], tenantY) + } +} + +// TestProcessFollowups_NilTeam_NoPanic verifies that when GetTeam returns (nil, nil) +// processFollowups skips without panicking, and no followup is sent. +func TestProcessFollowups_NilTeam_NoPanic(t *testing.T) { + tenantID := uuid.New() + teamID := uuid.New() + + ts := &stubTeamStore{ + followupTasks: []store.TeamTaskData{ + { + BaseModel: store.BaseModel{ID: uuid.New()}, + TeamID: teamID, + TenantID: tenantID, + FollowupChannel: "telegram", + FollowupChatID: "chat-1", + FollowupMessage: "ping", + }, + }, + getTeamFunc: func(_ context.Context, _ uuid.UUID) (*store.TeamData, error) { + return nil, nil // simulate nil without error + }, + } + as := &stubAgentStore{} + mb := bus.New() + tt := NewTaskTicker(ts, as, mb, 3600) + + // Must not panic. + tt.processFollowups(context.Background()) + + if ts.incrementCalls.Load() != 0 { + t.Errorf("IncrementFollowupCount = %d, want 0 (nil team skipped)", ts.incrementCalls.Load()) + } +} From 221cd78fcf1a97185fcd4095072d57e13a8a3b80 Mon Sep 17 00:00:00 2001 From: viettranx Date: Tue, 14 Apr 2026 10:09:45 +0700 Subject: [PATCH 2/3] fix(agent): emit tool.call event in parallel tool execution path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit v3 pipeline's parallel path (makeExecuteToolRaw) skipped the tool.call WebSocket event, so web UI and desktop UI silently dropped tool cards during real-time streaming. Only page refresh (which reloads history) revealed the tool calls. Both UIs rely on tool.call to create entries that later tool.result events can update. Fix: mirror the sequential path's emission in makeExecuteToolRaw. Bus.Broadcast is RWMutex-guarded, safe to call from parallel goroutines. Add tests at two layers to prevent regression: - Pipeline layer (stages_test.go): guards the dispatch contract — multiple tool calls route through ExecuteToolRaw + ProcessToolResult rather than ExecuteToolCall. Previously the parallel path had zero test coverage, which is why this bug escaped. - Agent layer (loop_pipeline_tool_callbacks_test.go): guards the emission contract — both sequential and parallel wrappers emit tool.call with correct payload and routing context. Mutation-verified. --- .../agent/loop_pipeline_tool_callbacks.go | 13 ++ .../loop_pipeline_tool_callbacks_test.go | 188 ++++++++++++++++++ internal/pipeline/stages_test.go | 61 ++++++ 3 files changed, 262 insertions(+) create mode 100644 internal/agent/loop_pipeline_tool_callbacks_test.go diff --git a/internal/agent/loop_pipeline_tool_callbacks.go b/internal/agent/loop_pipeline_tool_callbacks.go index 3012b5a6..b7743f3e 100644 --- a/internal/agent/loop_pipeline_tool_callbacks.go +++ b/internal/agent/loop_pipeline_tool_callbacks.go @@ -63,9 +63,22 @@ type toolRawResult struct { // makeExecuteToolRaw wraps tool I/O only (parallel-safe, no state mutation). // Returns tool message + toolRawResult (with timing + spanID) as opaque raw data for ProcessToolResult. func (l *Loop) makeExecuteToolRaw(req *RunRequest) func(ctx context.Context, tc providers.ToolCall) (providers.Message, any, error) { + emitRun := makeToolEmitRun(l, req) return func(ctx context.Context, tc providers.ToolCall) (providers.Message, any, error) { registryName := l.resolveToolCallName(tc.Name) argsJSON, _ := json.Marshal(tc.Arguments) + slog.Info("tool call", "agent", l.id, "tool", tc.Name, "args_len", len(argsJSON)) + + // Emit tool.call event at I/O start — parity with sequential path (makeExecuteToolCall). + // Without this, parallel tool execution (2+ concurrent tools) never notifies UI of + // tool invocation, so `tool.result` arrives with no matching `tool.call` to update. + // Bus.Broadcast is RWMutex-guarded; safe to call from parallel goroutines. + emitRun(AgentEvent{ + Type: protocol.AgentEventToolCall, + AgentID: l.id, + RunID: req.RunID, + Payload: map[string]any{"name": tc.Name, "id": tc.ID, "arguments": tc.Arguments}, + }) // Emit tool span start (goroutine-safe: channel send only). start := time.Now().UTC() diff --git a/internal/agent/loop_pipeline_tool_callbacks_test.go b/internal/agent/loop_pipeline_tool_callbacks_test.go new file mode 100644 index 00000000..866806c5 --- /dev/null +++ b/internal/agent/loop_pipeline_tool_callbacks_test.go @@ -0,0 +1,188 @@ +package agent + +import ( + "context" + "sync" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/pipeline" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/tools" + "github.com/nextlevelbuilder/goclaw/pkg/protocol" +) + +// stubExecutor implements tools.ToolExecutor with a canned successful Result. +// Used to isolate the tool-callback wrappers from real tool registry wiring. +type stubExecutor struct{} + +func (s *stubExecutor) ExecuteWithContext(_ context.Context, _ string, _ map[string]any, _, _, _, _ string, _ tools.AsyncCallback) *tools.Result { + return &tools.Result{ForLLM: "ok", IsError: false} +} +func (s *stubExecutor) TryActivateDeferred(string) bool { return false } +func (s *stubExecutor) ProviderDefs() []providers.ToolDefinition { return nil } +func (s *stubExecutor) Get(string) (tools.Tool, bool) { return nil, false } +func (s *stubExecutor) List() []string { return nil } +func (s *stubExecutor) Aliases() map[string]string { return nil } + +// eventCollector buffers AgentEvents for inspection in tests. +// Safe for concurrent appends from parallel goroutines. +type eventCollector struct { + mu sync.Mutex + events []AgentEvent +} + +func (c *eventCollector) onEvent(e AgentEvent) { + c.mu.Lock() + c.events = append(c.events, e) + c.mu.Unlock() +} + +func (c *eventCollector) filter(typ string) []AgentEvent { + c.mu.Lock() + defer c.mu.Unlock() + var out []AgentEvent + for _, e := range c.events { + if e.Type == typ { + out = append(out, e) + } + } + return out +} + +// newTestLoopForToolCallbacks builds a minimal Loop instance sufficient to +// exercise makeExecuteToolCall / makeExecuteToolRaw. All optional subsystems +// (tracing, metrics, input guard) are left nil and hit early-return paths. +func newTestLoopForToolCallbacks(onEvent func(AgentEvent)) *Loop { + return &Loop{ + id: "test-agent", + tools: &stubExecutor{}, + onEvent: onEvent, + } +} + +// TestMakeExecuteToolCall_EmitsToolCallEvent verifies the sequential wrapper +// emits a tool.call event before running tool I/O. +func TestMakeExecuteToolCall_EmitsToolCallEvent(t *testing.T) { + t.Parallel() + col := &eventCollector{} + l := newTestLoopForToolCallbacks(col.onEvent) + + req := &RunRequest{ + RunID: "run-1", + SessionKey: "sess-A", + UserID: "u-1", + Channel: "ws", + RunKind: "", + } + state := &pipeline.RunState{RunID: "run-1"} + tc := providers.ToolCall{ID: "tc-1", Name: "read_file", Arguments: map[string]any{"path": "/tmp/x"}} + + _, err := l.makeExecuteToolCall(req, &runState{})(context.Background(), state, tc) + if err != nil { + t.Fatalf("makeExecuteToolCall returned error: %v", err) + } + + calls := col.filter(protocol.AgentEventToolCall) + if len(calls) != 1 { + t.Fatalf("expected 1 tool.call event, got %d (all events: %+v)", len(calls), col.events) + } + assertToolCallPayload(t, calls[0], tc, req) +} + +// TestMakeExecuteToolRaw_EmitsToolCallEvent is the PRIMARY regression guard. +// The original bug: parallel path (makeExecuteToolRaw) did not emit tool.call, +// so web UI and desktop UI silently dropped tool info during real-time streaming. +// Mutation-verify: remove emitRun(...) from makeExecuteToolRaw — this test must fail. +func TestMakeExecuteToolRaw_EmitsToolCallEvent(t *testing.T) { + t.Parallel() + col := &eventCollector{} + l := newTestLoopForToolCallbacks(col.onEvent) + + req := &RunRequest{ + RunID: "run-2", + SessionKey: "sess-B", + UserID: "u-2", + Channel: "ws", + RunKind: "", + } + tc := providers.ToolCall{ID: "tc-2", Name: "write_file", Arguments: map[string]any{"path": "/tmp/y"}} + + msg, raw, err := l.makeExecuteToolRaw(req)(context.Background(), tc) + if err != nil { + t.Fatalf("makeExecuteToolRaw returned error: %v", err) + } + if msg.Role != "tool" || msg.ToolCallID != tc.ID { + t.Errorf("unexpected tool message: %+v", msg) + } + if raw == nil { + t.Error("expected non-nil raw data (toolRawResult)") + } + + calls := col.filter(protocol.AgentEventToolCall) + if len(calls) != 1 { + t.Fatalf("expected 1 tool.call event, got %d (all events: %+v)", len(calls), col.events) + } + assertToolCallPayload(t, calls[0], tc, req) +} + +// TestMakeExecuteToolRaw_ConcurrentCallsEmitAllEvents confirms the parallel +// wrapper is safe to invoke from multiple goroutines — mirrors the real +// executeParallel dispatch in pipeline/tool_stage.go. +func TestMakeExecuteToolRaw_ConcurrentCallsEmitAllEvents(t *testing.T) { + t.Parallel() + col := &eventCollector{} + l := newTestLoopForToolCallbacks(col.onEvent) + + req := &RunRequest{RunID: "run-3", SessionKey: "sess-C", UserID: "u-3", Channel: "ws"} + exec := l.makeExecuteToolRaw(req) + + const n = 5 + var wg sync.WaitGroup + for i := range n { + wg.Add(1) + go func(idx int) { + defer wg.Done() + tc := providers.ToolCall{ID: "tc-" + string(rune('a'+idx)), Name: "t", Arguments: nil} + if _, _, err := exec(context.Background(), tc); err != nil { + t.Errorf("goroutine %d: %v", idx, err) + } + }(i) + } + wg.Wait() + + calls := col.filter(protocol.AgentEventToolCall) + if len(calls) != n { + t.Fatalf("expected %d tool.call events, got %d", n, len(calls)) + } +} + +// assertToolCallPayload verifies the event carries the expected tc identity +// and routing context from RunRequest. +func assertToolCallPayload(t *testing.T, ev AgentEvent, tc providers.ToolCall, req *RunRequest) { + t.Helper() + if ev.AgentID != "test-agent" { + t.Errorf("AgentID: got %q, want test-agent", ev.AgentID) + } + if ev.RunID != req.RunID { + t.Errorf("RunID: got %q, want %q", ev.RunID, req.RunID) + } + if ev.SessionKey != req.SessionKey { + t.Errorf("SessionKey: got %q, want %q", ev.SessionKey, req.SessionKey) + } + if ev.Channel != req.Channel { + t.Errorf("Channel: got %q, want %q", ev.Channel, req.Channel) + } + if ev.UserID != req.UserID { + t.Errorf("UserID: got %q, want %q", ev.UserID, req.UserID) + } + payload, ok := ev.Payload.(map[string]any) + if !ok { + t.Fatalf("Payload is not map[string]any: %T", ev.Payload) + } + if payload["id"] != tc.ID { + t.Errorf("payload.id: got %v, want %q", payload["id"], tc.ID) + } + if payload["name"] != tc.Name { + t.Errorf("payload.name: got %v, want %q", payload["name"], tc.Name) + } +} diff --git a/internal/pipeline/stages_test.go b/internal/pipeline/stages_test.go index b778b1f1..7ab33a2a 100644 --- a/internal/pipeline/stages_test.go +++ b/internal/pipeline/stages_test.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "github.com/nextlevelbuilder/goclaw/internal/bootstrap" @@ -789,6 +790,66 @@ func TestToolStage_MultipleTools_Sequential_MessagesInOrder(t *testing.T) { } } +// Regression: v3 parallel path invokes ExecuteToolRaw + ProcessToolResult +// for every tool call. If this breaks, the `tool.call` WS event emitted +// inside makeExecuteToolRaw (loop_pipeline_tool_callbacks.go) stops firing +// and UIs go silent during real-time tool execution. +func TestToolStage_MultipleTools_ParallelPath_InvokesRawAndProcessForEach(t *testing.T) { + t.Parallel() + var rawMu, procMu sync.Mutex + rawCalls := []string{} + procCalls := []string{} + deps := &PipelineDeps{ + // Parallel path requires all three callbacks. ExecuteToolCall is required + // upfront (nil-guard) even though the parallel branch won't invoke it. + ExecuteToolCall: func(_ context.Context, _ *RunState, _ providers.ToolCall) ([]providers.Message, error) { + t.Fatal("ExecuteToolCall must NOT be called when parallel path is active") + return nil, nil + }, + ExecuteToolRaw: func(_ context.Context, tc providers.ToolCall) (providers.Message, any, error) { + rawMu.Lock() + rawCalls = append(rawCalls, tc.ID) + rawMu.Unlock() + return providers.Message{Role: "tool", Content: "raw:" + tc.Name, ToolCallID: tc.ID}, nil, nil + }, + ProcessToolResult: func(_ context.Context, _ *RunState, tc providers.ToolCall, rawMsg providers.Message, _ any) []providers.Message { + procMu.Lock() + procCalls = append(procCalls, tc.ID) + procMu.Unlock() + return []providers.Message{rawMsg} + }, + } + stage := NewToolStage(deps) + state := defaultState() + state.Think.LastResponse = &providers.ChatResponse{ + ToolCalls: []providers.ToolCall{ + {ID: "1", Name: "tool_a"}, + {ID: "2", Name: "tool_b"}, + {ID: "3", Name: "tool_c"}, + }, + } + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + + // Each tool must be passed through BOTH ExecuteToolRaw (where tool.call emits) + // AND ProcessToolResult (where tool.result emits). Missing either breaks UIs. + if len(rawCalls) != 3 { + t.Errorf("ExecuteToolRaw called %d times, want 3 (one per tool)", len(rawCalls)) + } + if len(procCalls) != 3 { + t.Errorf("ProcessToolResult called %d times, want 3", len(procCalls)) + } + // ProcessToolResult must run sequentially in original order (deterministic state mutation) + if len(procCalls) == 3 && (procCalls[0] != "1" || procCalls[1] != "2" || procCalls[2] != "3") { + t.Errorf("ProcessToolResult order = %v, want [1 2 3]", procCalls) + } + if state.Tool.TotalToolCalls != 3 { + t.Errorf("TotalToolCalls = %d, want 3", state.Tool.TotalToolCalls) + } +} + func TestToolStage_LoopKilled_ReturnsBreakLoop(t *testing.T) { t.Parallel() deps := &PipelineDeps{ From 677396c71d3aeded42c8e4d7734940644c1527f3 Mon Sep 17 00:00:00 2001 From: viettranx Date: Tue, 14 Apr 2026 10:46:53 +0700 Subject: [PATCH 3/3] fix(secure-cli): prevent deny_verbose from blocking --version via substring match matchesBinaryDeny used unanchored regex on joined args, causing `-v` pattern to false-positive on `--version`. Split deny_verbose into matchesBinaryVerbose with start-anchored per-arg matching: `-v` blocks `-v`, `-vv`, `-v=1` but not `--version`. deny_args keeps joined matching for multi-token patterns. --- internal/tools/credentialed_exec.go | 40 ++++++++++- internal/tools/credentialed_exec_test.go | 87 +++++++++++++++++++++++- 2 files changed, 124 insertions(+), 3 deletions(-) diff --git a/internal/tools/credentialed_exec.go b/internal/tools/credentialed_exec.go index aa5e47ab..74fd2fbf 100644 --- a/internal/tools/credentialed_exec.go +++ b/internal/tools/credentialed_exec.go @@ -140,6 +140,7 @@ func resolveAndMatchBinary(binaryName string, configPath *string) (string, error } // matchesBinaryDeny checks if the joined args string matches any per-binary deny pattern. +// Used for deny_args where patterns span multiple args (e.g. `auth\s+login`, `repo\s+delete`). // Returns the matched pattern string, or empty if allowed. func matchesBinaryDeny(args []string, denyPatternsJSON json.RawMessage) string { if len(denyPatternsJSON) == 0 { @@ -163,6 +164,40 @@ func matchesBinaryDeny(args []string, denyPatternsJSON json.RawMessage) string { return "" } +// matchesBinaryVerbose checks each arg token against verbose/debug flag patterns. +// Patterns are anchored at the START of each arg (but not the end), which allows: +// - `-v` to match `-v`, `-vv`, `-vvv` (verbosity escalation), `-v=1`, `-vq` (combined flags) +// - `--verbose` to match `--verbose`, `--verbose=true` +// - `-v` to NOT match `--version` (char 1 is `-`, not `v`) +// - `--verbose` to NOT match `--version` (diverges at char 5) +// +// This is intentional: verbose flags leak sensitive output (tokens in HTTP headers, +// API response bodies, OAuth flows). Start-anchored per-arg matching catches the +// real verbose family without false-positive on safe flags like `--version`. +// Returns the matched pattern string, or empty if allowed. +func matchesBinaryVerbose(args []string, denyPatternsJSON json.RawMessage) string { + if len(denyPatternsJSON) == 0 { + return "" + } + var patterns []string + if err := json.Unmarshal(denyPatternsJSON, &patterns); err != nil || len(patterns) == 0 { + return "" + } + for _, p := range patterns { + re, err := regexp.Compile("^(?:" + p + ")") + if err != nil { + slog.Warn("secure_cli.invalid_deny_pattern", "pattern", p, "error", err) + continue + } + for _, arg := range args { + if re.MatchString(arg) { + return p + } + } + } + return "" +} + // executeCredentialed runs a CLI command in Direct Exec Mode (no shell). // Credentials are injected as env vars into the child process only. // rawCommand is the original command string before shell-word parsing (preserves quoting). @@ -196,8 +231,9 @@ func (t *ExecTool) executeCredentialed(ctx context.Context, cred *store.SecureCL if p := matchesBinaryDeny(args, cred.DenyArgs); p != "" { return credentialedDenyError(binary, args, p) } - // Per-binary verbose deny check (deny_verbose) - if p := matchesBinaryDeny(args, cred.DenyVerbose); p != "" { + // Per-binary verbose deny check (deny_verbose) — per-arg start-anchored match + // so `-v` blocks `-v`/`-vv`/`-v=1` but not `--version`. + if p := matchesBinaryVerbose(args, cred.DenyVerbose); p != "" { return credentialedDenyError(binary, args, p) } diff --git a/internal/tools/credentialed_exec_test.go b/internal/tools/credentialed_exec_test.go index 6e950867..209f27af 100644 --- a/internal/tools/credentialed_exec_test.go +++ b/internal/tools/credentialed_exec_test.go @@ -1,6 +1,9 @@ package tools -import "testing" +import ( + "encoding/json" + "testing" +) func TestDetectShellOperators(t *testing.T) { tests := []struct { @@ -120,3 +123,85 @@ func TestParseCommandBinary(t *testing.T) { }) } } + +// TestMatchesBinaryVerbose verifies start-anchored per-arg matching for +// deny_verbose patterns. Regression guard: `-v` must NOT false-positive on +// `--version` (used by the system to probe CLI availability), but MUST still +// block real verbose flags (`-v`, `-vv`, `-v=1`, `--verbose=true`) to prevent +// leakage of tokens/request bodies via verbose output. +func TestMatchesBinaryVerbose(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{"--verbose", "-v"}) + gcloudPatterns, _ := json.Marshal([]string{"--verbosity=debug", "--log-http"}) + awsPatterns, _ := json.Marshal([]string{"--debug"}) + + tests := []struct { + name string + patterns json.RawMessage + args []string + wantHit bool + }{ + // --- regression: safe flags must pass --- + {"gh --version not blocked", ghPatterns, []string{"--version"}, false}, + {"gh version subcmd not blocked", ghPatterns, []string{"version"}, false}, + {"gh --help not blocked", ghPatterns, []string{"--help"}, false}, + {"gh api repos/x not blocked", ghPatterns, []string{"api", "repos/x"}, false}, + + // --- real verbose flags still blocked --- + {"gh -v blocked", ghPatterns, []string{"-v"}, true}, + {"gh --verbose blocked", ghPatterns, []string{"--verbose"}, true}, + {"gh -vv blocked (escalation)", ghPatterns, []string{"-vv"}, true}, + {"gh -vvv blocked (escalation)", ghPatterns, []string{"-vvv"}, true}, + {"gh --verbose=true blocked (equals form)", ghPatterns, []string{"--verbose=true"}, true}, + {"gh -v in middle of args blocked", ghPatterns, []string{"api", "-v", "repos/x"}, true}, + + // --- gcloud patterns: exact flag=value --- + {"gcloud --verbosity=debug blocked", gcloudPatterns, []string{"--verbosity=debug"}, true}, + {"gcloud --verbosity=info not blocked", gcloudPatterns, []string{"--verbosity=info"}, false}, + {"gcloud --log-http blocked", gcloudPatterns, []string{"--log-http"}, true}, + {"gcloud version not blocked", gcloudPatterns, []string{"version"}, false}, + + // --- aws --- + {"aws --debug blocked", awsPatterns, []string{"--debug"}, true}, + {"aws --debugger not blocked-worthy (prefix match)", awsPatterns, []string{"--debugger"}, true}, // acceptable: still debug family + {"aws --version not blocked", awsPatterns, []string{"--version"}, false}, + + // --- empty / no patterns --- + {"empty patterns", json.RawMessage(nil), []string{"--verbose"}, false}, + {"empty args", ghPatterns, []string{}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := matchesBinaryVerbose(tt.args, tt.patterns) + if (got != "") != tt.wantHit { + t.Errorf("matchesBinaryVerbose(%v) = %q, wantHit=%v", tt.args, got, tt.wantHit) + } + }) + } +} + +// TestMatchesBinaryDenyJoinedArgs verifies deny_args keeps joined-string +// matching so multi-token patterns like `auth\s+login` and `repo\s+delete` +// still work. +func TestMatchesBinaryDenyJoinedArgs(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`auth\s+`, `repo\s+delete`, `secret\s+`}) + + tests := []struct { + name string + args []string + wantHit bool + }{ + {"gh auth login blocked", []string{"auth", "login"}, true}, + {"gh repo delete blocked", []string{"repo", "delete", "foo/bar"}, true}, + {"gh secret set blocked", []string{"secret", "set", "TOKEN"}, true}, + {"gh api repos allowed", []string{"api", "repos/x"}, false}, + {"gh repo view allowed", []string{"repo", "view", "foo/bar"}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := matchesBinaryDeny(tt.args, ghPatterns) + if (got != "") != tt.wantHit { + t.Errorf("matchesBinaryDeny(%v) = %q, wantHit=%v", tt.args, got, tt.wantHit) + } + }) + } +}