diff --git a/cmd/gateway_consumer.go b/cmd/gateway_consumer.go index d501903e..48563014 100644 --- a/cmd/gateway_consumer.go +++ b/cmd/gateway_consumer.go @@ -79,6 +79,14 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents slog.Info("inbound debounce configured", "debounce_ms", debounceMs) + // Track background goroutines (subagent announces, teammate messages) + // so shutdown can wait for in-flight work to complete. + var bgWg sync.WaitGroup + defer func() { + bgWg.Wait() + slog.Info("inbound consumer: all background goroutines drained") + }() + for { msg, ok := msgBus.ConsumeInbound(ctx) if !ok { @@ -95,10 +103,10 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents } } - if handleSubagentAnnounce(ctx, msg, cfg, sched, channelMgr, msgBus, getAnnounceMu) { + if handleSubagentAnnounce(ctx, msg, cfg, sched, channelMgr, msgBus, getAnnounceMu, &bgWg) { continue } - if handleTeammateMessage(ctx, msg, cfg, sched, channelMgr, teamStore, agentStore, msgBus, postTurn, &taskRunSessions) { + if handleTeammateMessage(ctx, msg, cfg, sched, channelMgr, teamStore, agentStore, msgBus, postTurn, &taskRunSessions, &bgWg) { continue } if handleResetCommand(msg, cfg, sessStore) { diff --git a/cmd/gateway_consumer_handlers.go b/cmd/gateway_consumer_handlers.go index d51dbdaf..7480a2b9 100644 --- a/cmd/gateway_consumer_handlers.go +++ b/cmd/gateway_consumer_handlers.go @@ -17,6 +17,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/channels" "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/safego" "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/sessions" "github.com/nextlevelbuilder/goclaw/internal/store" @@ -35,6 +36,7 @@ func handleSubagentAnnounce( channelMgr *channels.Manager, msgBus *bus.MessageBus, getAnnounceMu func(string) *sync.Mutex, + bgWg *sync.WaitGroup, ) bool { if !(msg.Channel == tools.ChannelSystem && strings.HasPrefix(msg.SenderID, "subagent:")) { return false @@ -128,7 +130,10 @@ func handleSubagentAnnounce( // Handle announce asynchronously with per-session serialization. // The mutex ensures concurrent announces for the same session wait for // each other, so each reads up-to-date session history. + bgWg.Add(1) go func(sessionKey, origCh, chatID, senderID, label string, meta map[string]string, req agent.RunRequest) { + defer bgWg.Done() + defer safego.Recover(nil, "component", "subagent_announce", "session", sessionKey) mu := getAnnounceMu(sessionKey) mu.Lock() defer mu.Unlock() @@ -193,6 +198,7 @@ func handleTeammateMessage( msgBus *bus.MessageBus, postTurn tools.PostTurnProcessor, taskRunSessions *sync.Map, + bgWg *sync.WaitGroup, ) bool { if !(msg.Channel == tools.ChannelSystem && strings.HasPrefix(msg.SenderID, "teammate:")) { return false @@ -286,7 +292,10 @@ func handleTeammateMessage( LinkedTraceID: linkedTraceID, }) + bgWg.Add(1) go func(origCh, origChatID, senderID, taskID string, meta, inMeta map[string]string) { + defer bgWg.Done() + defer safego.Recover(nil, "component", "teammate_message", "task_id", taskID) // Lock renewal heartbeat: extend task lock every 10 min to prevent // the ticker from recovering long-running tasks as stale. var lockStop chan struct{} diff --git a/internal/channels/slack/channel.go b/internal/channels/slack/channel.go index 61538cae..ce44b5f0 100644 --- a/internal/channels/slack/channel.go +++ b/internal/channels/slack/channel.go @@ -15,6 +15,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/channels" "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/safego" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -167,12 +168,14 @@ func (c *Channel) Start(ctx context.Context) error { // Goroutine 1: Event loop go func() { defer c.wg.Done() + defer safego.Recover(nil, "component", "slack_event_loop") c.eventLoop(smCtx) }() // Goroutine 2: Socket Mode connection with dead socket error classification go func() { defer c.wg.Done() + defer safego.Recover(nil, "component", "slack_socket_mode") for { if err := c.sm.RunContext(smCtx); err != nil { if smCtx.Err() != nil { @@ -192,6 +195,7 @@ func (c *Channel) Start(ctx context.Context) error { // Goroutine 3: Periodic sweep (every 2 minutes) for TTL-based map eviction go func() { defer c.wg.Done() + defer safego.Recover(nil, "component", "slack_sweep") ticker := time.NewTicker(2 * time.Minute) defer ticker.Stop() for { diff --git a/internal/http/skills.go b/internal/http/skills.go index 63204732..87b14427 100644 --- a/internal/http/skills.go +++ b/internal/http/skills.go @@ -203,7 +203,7 @@ func (h *SkillsHandler) handleInstallDeps(w http.ResponseWriter, r *http.Request if !h.requireMasterTenant(w, r) { return } - dirs := h.skills.ListSystemSkillDirs() + dirs := h.skills.ListSystemSkillDirs(r.Context()) if len(dirs) == 0 { writeJSON(w, http.StatusOK, map[string]string{"message": "no system skills"}) return @@ -342,7 +342,7 @@ func (h *SkillsHandler) rescanAndUpdate() (updated int, results []depResult) { if manifest == nil || manifest.IsEmpty() { // No deps needed — if archived, recover to active and clear stale deps. if sk.Status == "archived" { - _ = h.skills.StoreMissingDeps(id, nil) + _ = h.skills.StoreMissingDeps(masterCtx, id, nil) _ = h.skills.UpdateSkill(masterCtx, id, map[string]any{"status": "active"}) results = append(results, depResult{Slug: sk.Slug, Status: "active"}) updated++ @@ -354,7 +354,7 @@ func (h *SkillsHandler) rescanAndUpdate() (updated int, results []depResult) { } ok, missing := skills.CheckSkillDeps(manifest) - _ = h.skills.StoreMissingDeps(id, missing) + _ = h.skills.StoreMissingDeps(masterCtx, id, missing) switch { case ok && sk.Status == "archived": @@ -471,7 +471,7 @@ func (h *SkillsHandler) handleToggle(w http.ResponseWriter, r *http.Request) { manifest := h.scanWithFallback(sk) if manifest != nil && !manifest.IsEmpty() { depOk, missing := skills.CheckSkillDeps(manifest) - _ = h.skills.StoreMissingDeps(id, missing) + _ = h.skills.StoreMissingDeps(r.Context(), id, missing) if depOk { newStatus = "active" } else { diff --git a/internal/skills/seeder.go b/internal/skills/seeder.go index a8d71fc5..057d371d 100644 --- a/internal/skills/seeder.go +++ b/internal/skills/seeder.go @@ -22,7 +22,7 @@ type SystemSkillStore interface { GetNextVersion(slug string) int BumpVersion() UpdateSkill(ctx context.Context, id uuid.UUID, updates map[string]interface{}) error - StoreMissingDeps(id uuid.UUID, missing []string) error + StoreMissingDeps(ctx context.Context, id uuid.UUID, missing []string) error } // seededSkill tracks a skill that was seeded and needs async dep checking. @@ -173,7 +173,7 @@ func (s *Seeder) CheckDepsAsync(skills []seededSkill, msgBus *bus.MessageBus) { ok, missing := CheckSkillDeps(manifest) // Always persist missing deps so UI can display them per-skill - _ = s.store.StoreMissingDeps(sk.id, missing) + _ = s.store.StoreMissingDeps(context.Background(), sk.id, missing) status := "active" if !ok { status = "archived" diff --git a/internal/store/pg/cron.go b/internal/store/pg/cron.go index 882e9ece..14a21259 100644 --- a/internal/store/pg/cron.go +++ b/internal/store/pg/cron.go @@ -1,6 +1,7 @@ package pg import ( + "context" "database/sql" "log/slog" "sync" @@ -15,12 +16,14 @@ const defaultCronCacheTTL = 2 * time.Minute // PGCronStore implements store.CronStore backed by Postgres. // GetDueJobs() uses an in-memory cache with TTL to reduce DB polling (1s interval). type PGCronStore struct { - db *sql.DB - mu sync.Mutex - onJob func(job *store.CronJob) (*store.CronJobResult, error) - onEvent func(event store.CronEvent) - running bool - stop chan struct{} + db *sql.DB + mu sync.Mutex + baseCtx context.Context // lifecycle context for background DB operations + cancelCtx context.CancelFunc // cancelled on Stop() + onJob func(job *store.CronJob) (*store.CronJobResult, error) + onEvent func(event store.CronEvent) + running bool + stop chan struct{} // Job cache: reduces GetDueJobs polling from 86,400 queries/day to ~720/day jobCache []store.CronJob @@ -57,6 +60,7 @@ func (s *PGCronStore) Start() error { if s.running { return nil } + s.baseCtx, s.cancelCtx = context.WithCancel(context.Background()) s.stop = make(chan struct{}) s.running = true s.recomputeStaleJobs() @@ -72,6 +76,9 @@ func (s *PGCronStore) Stop() { return } close(s.stop) + if s.cancelCtx != nil { + s.cancelCtx() + } s.running = false } diff --git a/internal/store/pg/cron_crud.go b/internal/store/pg/cron_crud.go index bebe05e1..1eef96ad 100644 --- a/internal/store/pg/cron_crud.go +++ b/internal/store/pg/cron_crud.go @@ -66,7 +66,7 @@ func (s *PGCronStore) AddJob(ctx context.Context, name string, schedule store.Cr nextRun := computeNextRun(&schedule, now, s.defaultTZ) - _, err := s.db.Exec( + _, err := s.db.ExecContext(ctx, `INSERT INTO cron_jobs (id, tenant_id, agent_id, user_id, name, enabled, schedule_kind, cron_expression, run_at, timezone, interval_ms, payload, delete_after_run, next_run_at, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, true, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)`, @@ -133,7 +133,7 @@ func (s *PGCronStore) ListJobs(ctx context.Context, includeDisabled bool, agentI q += " ORDER BY created_at DESC" - rows, err := s.db.Query(q, args...) + rows, err := s.db.QueryContext(ctx, q, args...) if err != nil { return nil } diff --git a/internal/store/pg/cron_exec.go b/internal/store/pg/cron_exec.go index 0c8da01c..514090c1 100644 --- a/internal/store/pg/cron_exec.go +++ b/internal/store/pg/cron_exec.go @@ -27,7 +27,7 @@ func (s *PGCronStore) RunJob(ctx context.Context, jobID string, force bool) (boo // Mark job as running before execution if id, parseErr := uuid.Parse(jobID); parseErr == nil { - s.db.Exec("UPDATE cron_jobs SET last_status = 'running', updated_at = $1 WHERE id = $2", time.Now(), id) + s.db.ExecContext(ctx, "UPDATE cron_jobs SET last_status = 'running', updated_at = $1 WHERE id = $2", time.Now(), id) } s.mu.Lock() s.cacheLoaded = false diff --git a/internal/store/pg/cron_scheduler.go b/internal/store/pg/cron_scheduler.go index e8c97c0b..995ac7d6 100644 --- a/internal/store/pg/cron_scheduler.go +++ b/internal/store/pg/cron_scheduler.go @@ -33,7 +33,7 @@ func (s *PGCronStore) GetDueJobs(now time.Time) []store.CronJob { // refreshJobCache reloads all enabled jobs from DB. Must be called with mu held. func (s *PGCronStore) refreshJobCache() { - rows, err := s.db.Query( + rows, err := s.db.QueryContext(s.baseCtx, `SELECT id, tenant_id, agent_id, user_id, name, enabled, schedule_kind, cron_expression, run_at, timezone, interval_ms, payload, delete_after_run, next_run_at, last_run_at, last_status, last_error, created_at, updated_at FROM cron_jobs WHERE enabled = true`) @@ -65,7 +65,7 @@ func (s *PGCronStore) InvalidateCache() { // This happens when the gateway was stopped/crashed while a job was executing, // or when the previously computed next_run_at was consumed but never recomputed. func (s *PGCronStore) recomputeStaleJobs() { - rows, err := s.db.Query( + rows, err := s.db.QueryContext(s.baseCtx, `SELECT id, schedule_kind, cron_expression, run_at, timezone, interval_ms FROM cron_jobs WHERE enabled = true AND next_run_at IS NULL`) if err != nil { @@ -105,12 +105,12 @@ func (s *PGCronStore) recomputeStaleJobs() { next := computeNextRun(&schedule, now, s.defaultTZ) if next == nil { if scheduleKind == "at" { - s.db.Exec("UPDATE cron_jobs SET enabled = false, updated_at = $1 WHERE id = $2", now, id) + s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET enabled = false, updated_at = $1 WHERE id = $2", now, id) } continue } - s.db.Exec("UPDATE cron_jobs SET next_run_at = $1, updated_at = $2 WHERE id = $3", *next, now, id) + s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET next_run_at = $1, updated_at = $2 WHERE id = $3", *next, now, id) fixed++ } @@ -149,7 +149,7 @@ func (s *PGCronStore) checkAndRunDueJobs() { // Clear next_run for all due jobs first to prevent duplicate fires for _, job := range dueJobs { if id, parseErr := uuid.Parse(job.ID); parseErr == nil { - s.db.Exec("UPDATE cron_jobs SET next_run_at = NULL WHERE id = $1", id) + s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET next_run_at = NULL WHERE id = $1", id) } } @@ -222,7 +222,7 @@ func (s *PGCronStore) executeOneJob(job store.CronJob, handler func(job *store.C if aid, aidErr := uuid.Parse(job.AgentID); aidErr == nil { agentUUID = &aid } - s.db.Exec( + s.db.ExecContext(s.baseCtx, `INSERT INTO cron_run_logs (id, job_id, agent_id, status, error, summary, duration_ms, input_tokens, output_tokens, ran_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, logID, id, agentUUID, status, lastError, summary, durationMS, inputTokens, outputTokens, now, @@ -232,12 +232,12 @@ func (s *PGCronStore) executeOneJob(job store.CronJob, handler func(job *store.C // Recompute next run or delete if job.DeleteAfterRun { if id, parseErr := uuid.Parse(job.ID); parseErr == nil { - s.db.Exec("DELETE FROM cron_jobs WHERE id = $1", id) + s.db.ExecContext(s.baseCtx, "DELETE FROM cron_jobs WHERE id = $1", id) } } else if id, parseErr := uuid.Parse(job.ID); parseErr == nil { schedule := job.Schedule next := computeNextRun(&schedule, now, s.defaultTZ) - s.db.Exec( + s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET last_run_at = $1, last_status = $2, last_error = $3, next_run_at = $4, updated_at = $5 WHERE id = $6", now, status, lastError, next, now, id, ) diff --git a/internal/store/pg/sessions.go b/internal/store/pg/sessions.go index b325a9df..e9415ddd 100644 --- a/internal/store/pg/sessions.go +++ b/internal/store/pg/sessions.go @@ -41,7 +41,7 @@ func NewPGSessionStore(db *sql.DB) *PGSessionStore { // The last hyphen-delimited segment is the base36 timestamp used as convId. // Idempotent — no-op if no legacy keys exist. func (s *PGSessionStore) migrateLegacyWSKeys() { - res, err := s.db.Exec(` + res, err := s.db.ExecContext(context.Background(), ` UPDATE sessions SET session_key = regexp_replace( session_key, @@ -103,7 +103,7 @@ func (s *PGSessionStore) GetOrCreate(ctx context.Context, key string) *store.Ses s.cache[sessionCacheKey(ctx, key)] = data msgsJSON, _ := json.Marshal([]providers.Message{}) - s.db.Exec( + s.db.ExecContext(ctx, `INSERT INTO sessions (id, session_key, messages, created_at, updated_at, team_id, tenant_id) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (tenant_id, session_key) DO NOTHING`, uuid.Must(uuid.NewV7()), key, msgsJSON, now, now, teamID, tenantIDForInsert(ctx), diff --git a/internal/store/pg/sessions_list.go b/internal/store/pg/sessions_list.go index 9e6192cc..2d574d30 100644 --- a/internal/store/pg/sessions_list.go +++ b/internal/store/pg/sessions_list.go @@ -352,7 +352,7 @@ func (s *PGSessionStore) getOrInit(ctx context.Context, key string) *store.Sessi s.cache[sessionCacheKey(ctx, key)] = data msgsJSON, _ := json.Marshal([]providers.Message{}) - s.db.Exec( + s.db.ExecContext(ctx, `INSERT INTO sessions (id, session_key, messages, created_at, updated_at, tenant_id) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (tenant_id, session_key) DO NOTHING`, uuid.Must(uuid.NewV7()), key, msgsJSON, now, now, tenantIDForInsert(ctx), diff --git a/internal/store/pg/skills.go b/internal/store/pg/skills.go index b406ddf5..b262cd81 100644 --- a/internal/store/pg/skills.go +++ b/internal/store/pg/skills.go @@ -188,7 +188,7 @@ func (s *PGSkillStore) scanSkillInfoList(rows *sql.Rows) []store.SkillInfo { // StoreMissingDeps persists the missing_deps list for a skill into the deps JSONB column. // Only updates system skills unscoped; custom skills require tenant match. -func (s *PGSkillStore) StoreMissingDeps(id uuid.UUID, missing []string) error { +func (s *PGSkillStore) StoreMissingDeps(ctx context.Context, id uuid.UUID, missing []string) error { if missing == nil { missing = []string{} } @@ -198,7 +198,7 @@ func (s *PGSkillStore) StoreMissingDeps(id uuid.UUID, missing []string) error { } // System skills can be updated without tenant; custom skills need tenant scope. // Use is_system check to ensure cross-tenant safety for custom skills. - _, err = s.db.Exec( + _, err = s.db.ExecContext(ctx, `UPDATE skills SET deps = $1, updated_at = NOW() WHERE id = $2 AND is_system = true`, encoded, id, ) diff --git a/internal/store/pg/skills_admin.go b/internal/store/pg/skills_admin.go index 7eb18779..ee5ee871 100644 --- a/internal/store/pg/skills_admin.go +++ b/internal/store/pg/skills_admin.go @@ -77,8 +77,8 @@ func (s *PGSkillStore) UpsertSystemSkill(ctx context.Context, p SkillCreateParam // ListSystemSkillDirs returns slug->file_path map for all enabled system skills. // Disabled system skills are excluded — dep checking and injection are skipped for them. -func (s *PGSkillStore) ListSystemSkillDirs() map[string]string { - rows, err := s.db.Query( +func (s *PGSkillStore) ListSystemSkillDirs(ctx context.Context) map[string]string { + rows, err := s.db.QueryContext(ctx, `SELECT slug, file_path FROM skills WHERE is_system = true AND enabled = true`) if err != nil { return nil diff --git a/internal/tools/publish_skill.go b/internal/tools/publish_skill.go index c85f1bd9..1a477175 100644 --- a/internal/tools/publish_skill.go +++ b/internal/tools/publish_skill.go @@ -175,7 +175,7 @@ func (t *PublishSkillTool) Execute(ctx context.Context, args map[string]any) *Re if manifest != nil && !manifest.IsEmpty() { ok, missing := skills.CheckSkillDeps(manifest) if !ok { - _ = t.skills.StoreMissingDeps(id, missing) + _ = t.skills.StoreMissingDeps(ctx, id, missing) depsWarning = skills.FormatMissing(missing) } } diff --git a/internal/tools/skill_manage.go b/internal/tools/skill_manage.go index 7f9e8eec..f3b7101a 100644 --- a/internal/tools/skill_manage.go +++ b/internal/tools/skill_manage.go @@ -198,7 +198,7 @@ func (t *SkillManageTool) executeCreate(ctx context.Context, args map[string]any if manifest != nil && !manifest.IsEmpty() { ok, missing := skills.CheckSkillDeps(manifest) if !ok { - _ = t.skills.StoreMissingDeps(id, missing) + _ = t.skills.StoreMissingDeps(ctx, id, missing) depsWarning = skills.FormatMissing(missing) } }